Implement State DB Methods (#3193)

* Added state implementation

* Gaze

* Fixed test

* Fixed build file

* Fixed all tests

* Merged with master

* Added comments to save and get from roots

* Make it explicit signing root

* s/./,

* s/marshalled/marshaled

beacon-chain/db/db.go

@@ -39,7 +39,7 @@ type Database interface {
 	ValidatorLatestVote(ctx context.Context, validatorIdx uint64) (*pb.ValidatorLatestVote, error)
 	HasValidatorLatestVote(ctx context.Context, validatorIdx uint64) bool
 	SaveValidatorLatestVote(ctx context.Context, validatorIdx uint64, vote *pb.ValidatorLatestVote) error
-	State(ctx context.Context, f *filters.QueryFilter) (*pb.BeaconState, error)
+	State(ctx context.Context, blockRoot [32]byte) (*pb.BeaconState, error)
 	HeadState(ctx context.Context) (*pb.BeaconState, error)
 	SaveState(ctx context.Context, state *pb.BeaconState, blockRoot [32]byte) error
 	ValidatorIndex(ctx context.Context, publicKey [48]byte) (uint64, bool, error)

beacon-chain/db/kv/BUILD.bazel

@@ -29,6 +29,7 @@ go_test(
         "attestations_test.go",
         "blocks_test.go",
         "kv_test.go",
+        "state_test.go",
         "validators_test.go",
     ],
     embed = [":go_default_library"],

beacon-chain/db/kv/state.go

@@ -3,24 +3,73 @@ package kv
 import (
 	"context"
 
-	"github.com/prysmaticlabs/prysm/beacon-chain/db/filters"
+	"github.com/boltdb/bolt"
+	"github.com/gogo/protobuf/proto"
+	"github.com/pkg/errors"
 	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
 )
 
-// State retrieval by filter criteria.
-// TODO(#3164): Implement.
-func (k *Store) State(ctx context.Context, f *filters.QueryFilter) (*pb.BeaconState, error) {
-	return nil, nil
+// State returns the saved state using block's signing root,
+// this particular block was used to generate the state.
+func (k *Store) State(ctx context.Context, blockRoot [32]byte) (*pb.BeaconState, error) {
+	var s *pb.BeaconState
+	err := k.db.View(func(tx *bolt.Tx) error {
+		bucket := tx.Bucket(stateBucket)
+		enc := bucket.Get(blockRoot[:])
+		if enc == nil {
+			return nil
+		}
+
+		var err error
+		s, err = createState(enc)
+
+		return err
+	})
+	return s, err
 }
 
-// HeadState returns the latest canonical state in eth2.
-// TODO(#3164): Implement.
+// HeadState returns the latest canonical state in beacon chain.
 func (k *Store) HeadState(ctx context.Context) (*pb.BeaconState, error) {
-	return nil, nil
+	var s *pb.BeaconState
+	err := k.db.View(func(tx *bolt.Tx) error {
+		// Retrieve head block's signing root from blocks bucket,
+		// to look up what the head state is.
+		bucket := tx.Bucket(blocksBucket)
+		headBlkRoot := bucket.Get(headBlockRootKey)
+
+		bucket = tx.Bucket(stateBucket)
+		enc := bucket.Get(headBlkRoot)
+		if enc == nil {
+			return nil
+		}
+
+		var err error
+		s, err = createState(enc)
+
+		return err
+	})
+	return s, err
 }
 
-// SaveState stores a state to the db by the block root which triggered it.
-// TODO(#3164): Implement.
+// SaveState stores a state to the db using block's signing root which was used to generate the state.
 func (k *Store) SaveState(ctx context.Context, state *pb.BeaconState, blockRoot [32]byte) error {
-	return nil
+	enc, err := proto.Marshal(state)
+	if err != nil {
+		return err
+	}
+
+	return k.db.Update(func(tx *bolt.Tx) error {
+		bucket := tx.Bucket(stateBucket)
+		return bucket.Put(blockRoot[:], enc)
+	})
+}
+
+// creates state from marshaled proto state bytes.
+func createState(enc []byte) (*pb.BeaconState, error) {
+	protoState := &pb.BeaconState{}
+	err := proto.Unmarshal(enc, protoState)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to unmarshal encoding")
+	}
+	return protoState, nil
 }

beacon-chain/db/kv/state_test.go

@@ -0,0 +1,77 @@
+package kv
+
+import (
+	"context"
+	"reflect"
+	"testing"
+
+	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
+)
+
+func TestState_CanSaveRetrieve(t *testing.T) {
+	db := setupDB(t)
+	defer teardownDB(t, db)
+
+	s := &pb.BeaconState{Slot: 100}
+	r := [32]byte{'A'}
+
+	if err := db.SaveState(context.Background(), s, r); err != nil {
+		t.Fatal(err)
+	}
+
+	savedS, err := db.State(context.Background(), r)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if !reflect.DeepEqual(s, savedS) {
+		t.Error("did not retrieve saved state")
+	}
+
+	savedS, err = db.State(context.Background(), [32]byte{'B'})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if savedS != nil {
+		t.Error("unsaved state should've been nil")
+	}
+}
+
+func TestHeadState_CanSaveRetrieve(t *testing.T) {
+	db := setupDB(t)
+	defer teardownDB(t, db)
+
+	s := &pb.BeaconState{Slot: 100}
+	headRoot := [32]byte{'A'}
+
+	if err := db.SaveHeadBlockRoot(context.Background(), headRoot); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := db.SaveState(context.Background(), s, headRoot); err != nil {
+		t.Fatal(err)
+	}
+
+	savedHeadS, err := db.HeadState(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if !reflect.DeepEqual(s, savedHeadS) {
+		t.Error("did not retrieve saved state")
+	}
+
+	if err := db.SaveHeadBlockRoot(context.Background(), [32]byte{'B'}); err != nil {
+		t.Fatal(err)
+	}
+
+	savedHeadS, err = db.HeadState(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if savedHeadS != nil {
+		t.Error("unsaved head state should've been nil")
+	}
+}