From 15f08d6aade01b24049d90d2c4c310960438d355 Mon Sep 17 00:00:00 2001
From: Eugene Kim <ek@harmony.one>
Date: Mon, 2 Dec 2019 14:11:26 -0800
Subject: [PATCH] =?UTF-8?q?ChainConfigForNetworkType=20=E2=86=92=20Network?=
 =?UTF-8?q?Type.ChainConfig?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Makes the call sites read more naturally.
---
 cmd/harmony/main.go             |  2 +-
 internal/configs/node/config.go | 13 +++++++++++++
 node/node.go                    | 13 +------------
 3 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/cmd/harmony/main.go b/cmd/harmony/main.go
index bde8b99f3..96da9194d 100644
--- a/cmd/harmony/main.go
+++ b/cmd/harmony/main.go
@@ -459,7 +459,7 @@ func main() {
 				os.Exit(1)
 			}
 			// Use the number of shards as of staking epoch.
-			chainConfig := node.ChainConfigForNetworkType(nodeconfig.NetworkType(*networkType))
+			chainConfig := nodeconfig.NetworkType(*networkType).ChainConfig()
 			stakingEpochShardConfig := shard.Schedule.InstanceForEpoch(chainConfig.StakingEpoch)
 			numShardsBig := big.NewInt(int64(stakingEpochShardConfig.NumShards()))
 			shardIDBig := new(big.Int).Mod(blsPubKey.Big(), numShardsBig)
diff --git a/internal/configs/node/config.go b/internal/configs/node/config.go
index e770f8d85..cf487cec4 100644
--- a/internal/configs/node/config.go
+++ b/internal/configs/node/config.go
@@ -10,6 +10,8 @@ import (
 
 	"github.com/harmony-one/bls/ffi/go/bls"
 	p2p_crypto "github.com/libp2p/go-libp2p-crypto"
+
+	"github.com/harmony-one/harmony/internal/params"
 )
 
 // Role defines a role of a node.
@@ -241,3 +243,14 @@ func SetPublicRPC(v bool) {
 func GetPublicRPC() bool {
 	return publicRPC
 }
+
+func (t NetworkType) ChainConfig() params.ChainConfig {
+	switch t {
+	case Mainnet:
+		return *params.MainnetChainConfig
+	case Pangaea:
+		return *params.PangaeaChainConfig
+	default:
+		return *params.TestnetChainConfig
+	}
+}
diff --git a/node/node.go b/node/node.go
index c24a6c924..8c584f4dc 100644
--- a/node/node.go
+++ b/node/node.go
@@ -408,7 +408,7 @@ func New(host p2p.Host, consensusObj *consensus.Consensus,
 	}
 
 	networkType := node.NodeConfig.GetNetworkType()
-	chainConfig := ChainConfigForNetworkType(networkType)
+	chainConfig := networkType.ChainConfig()
 	node.chainConfig = chainConfig
 
 	collection := shardchain.NewCollection(
@@ -483,17 +483,6 @@ func New(host p2p.Host, consensusObj *consensus.Consensus,
 	return &node
 }
 
-func ChainConfigForNetworkType(t nodeconfig.NetworkType) params.ChainConfig {
-	switch t {
-	case nodeconfig.Mainnet:
-		return *params.MainnetChainConfig
-	case nodeconfig.Pangaea:
-		return *params.PangaeaChainConfig
-	default:
-		return *params.TestnetChainConfig
-	}
-}
-
 // InitConsensusWithValidators initialize shard state from latest epoch and update committee pub
 // keys for consensus and drand
 func (node *Node) InitConsensusWithValidators() (err error) {