From 84a2582877b941b4b03c89a467a4741ef23a671f Mon Sep 17 00:00:00 2001
From: colin axner <colinaxner@berkeley.edu>
Date: Tue, 18 Jun 2019 15:11:31 -0700
Subject: [PATCH] remove BlockHeight in context (#4580)

closes #4578
---
 .../improvements/sdk/4580-remove-block-he     |  1 +
 types/context.go                              |  6 ++---
 types/context_test.go                         | 27 +++++++++++++++++++
 3 files changed, 30 insertions(+), 4 deletions(-)
 create mode 100644 .pending/improvements/sdk/4580-remove-block-he

diff --git a/.pending/improvements/sdk/4580-remove-block-he b/.pending/improvements/sdk/4580-remove-block-he
new file mode 100644
index 000000000000..c7901c036a7b
--- /dev/null
+++ b/.pending/improvements/sdk/4580-remove-block-he
@@ -0,0 +1 @@
+#4580 Update `Context#BlockHeight` to properly set the block height via `WithBlockHeader`.
diff --git a/types/context.go b/types/context.go
index f626fc40caa6..de426bfe7e4a 100644
--- a/types/context.go
+++ b/types/context.go
@@ -43,7 +43,6 @@ func NewContext(ms MultiStore, header abci.Header, isCheckTx bool, logger log.Lo
 	}
 	c = c.WithMultiStore(ms)
 	c = c.WithBlockHeader(header)
-	c = c.WithBlockHeight(header.Height)
 	c = c.WithChainID(header.ChainID)
 	c = c.WithIsCheckTx(isCheckTx)
 	c = c.WithTxBytes(nil)
@@ -136,7 +135,6 @@ type contextKey int // local to the context module
 const (
 	contextKeyMultiStore contextKey = iota
 	contextKeyBlockHeader
-	contextKeyBlockHeight
 	contextKeyChainID
 	contextKeyIsCheckTx
 	contextKeyTxBytes
@@ -154,7 +152,7 @@ func (c Context) MultiStore() MultiStore {
 
 func (c Context) BlockHeader() abci.Header { return c.Value(contextKeyBlockHeader).(abci.Header) }
 
-func (c Context) BlockHeight() int64 { return c.Value(contextKeyBlockHeight).(int64) }
+func (c Context) BlockHeight() int64 { return c.BlockHeader().Height }
 
 func (c Context) ChainID() string { return c.Value(contextKeyChainID).(string) }
 
@@ -202,7 +200,7 @@ func (c Context) WithProposer(addr ConsAddress) Context {
 func (c Context) WithBlockHeight(height int64) Context {
 	newHeader := c.BlockHeader()
 	newHeader.Height = height
-	return c.withValue(contextKeyBlockHeight, height).withValue(contextKeyBlockHeader, newHeader)
+	return c.WithBlockHeader(newHeader)
 }
 
 func (c Context) WithChainID(chainID string) Context { return c.withValue(contextKeyChainID, chainID) }
diff --git a/types/context_test.go b/types/context_test.go
index 56f951584e45..bfa66e0745f8 100644
--- a/types/context_test.go
+++ b/types/context_test.go
@@ -2,6 +2,7 @@ package types_test
 
 import (
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 
@@ -10,6 +11,8 @@ import (
 
 	abci "github.com/tendermint/tendermint/abci/types"
 
+	"github.com/tendermint/tendermint/crypto/secp256k1"
+
 	"github.com/cosmos/cosmos-sdk/store"
 	"github.com/cosmos/cosmos-sdk/types"
 )
@@ -184,3 +187,27 @@ func TestContextWithCustom(t *testing.T) {
 	require.Equal(t, meter, ctx.GasMeter())
 	require.Equal(t, minGasPrices, ctx.MinGasPrices())
 }
+
+// Testing saving/loading of header fields to/from the context
+func TestContextHeader(t *testing.T) {
+	var ctx types.Context
+
+	require.Panics(t, func() { ctx.BlockHeader() })
+	require.Panics(t, func() { ctx.BlockHeight() })
+
+	height := int64(5)
+	time := time.Now()
+	addr := secp256k1.GenPrivKey().PubKey().Address()
+	proposer := types.ConsAddress(addr)
+
+	ctx = types.NewContext(nil, abci.Header{}, false, nil)
+
+	ctx = ctx.
+		WithBlockHeight(height).
+		WithBlockTime(time).
+		WithProposer(proposer)
+	require.Equal(t, height, ctx.BlockHeight())
+	require.Equal(t, height, ctx.BlockHeader().Height)
+	require.Equal(t, time, ctx.BlockHeader().Time)
+	require.Equal(t, proposer.Bytes(), ctx.BlockHeader().ProposerAddress)
+}