optimise

[bytom/vapor.git] / wallet / annotated.go

diff --git a/wallet/annotated.go b/wallet/annotated.go
index ef8cfc3..2fb1adc 100644 (file)
--- a/wallet/annotated.go
+++ b/wallet/annotated.go
@@ -3,6 +3,8 @@ package wallet
 import (
        "encoding/json"
 
+       btmCommon "github.com/bytom/common"
+       btmConsensus "github.com/bytom/consensus"
        log "github.com/sirupsen/logrus"
 
        "github.com/vapor/account"
@@ -190,9 +192,10 @@ func (w *Wallet) BuildAnnotatedInput(tx *types.Tx, i uint32) *query.AnnotatedInp
 }
 
 func (w *Wallet) getAddressFromControlProgram(prog []byte, isMainchain bool) string {
-       netParams := &consensus.ActiveNetParams
+       var netParams interface{}
+       netParams = &consensus.ActiveNetParams
        if isMainchain {
-               netParams = &consensus.MainNetParams
+               netParams = &btmConsensus.MainNetParams
        }
 
        if segwit.IsP2WPKHScript(prog) {
@@ -208,22 +211,44 @@ func (w *Wallet) getAddressFromControlProgram(prog []byte, isMainchain bool) str
        return ""
 }
 
-func BuildP2PKHAddress(pubHash []byte, netParams *consensus.Params) string {
-       address, err := common.NewAddressWitnessPubKeyHash(pubHash, netParams)
-       if err != nil {
-               return ""
-       }
+// BuildP2PKHAddress create P2PKH address
+func BuildP2PKHAddress(pubHash []byte, netParams interface{}) string {
+       switch netParams.(type) {
+       case *consensus.Params:
+               params := netParams.(*consensus.Params)
+               address, err := common.NewAddressWitnessPubKeyHash(pubHash, params)
+               if err == nil {
+                       return address.EncodeAddress()
+               }
 
-       return address.EncodeAddress()
+       case *btmConsensus.Params:
+               params := netParams.(*btmConsensus.Params)
+               address, err := btmCommon.NewAddressWitnessPubKeyHash(pubHash, params)
+               if err == nil {
+                       return address.EncodeAddress()
+               }
+       }
+       return ""
 }
 
-func BuildP2SHAddress(scriptHash []byte, netParams *consensus.Params) string {
-       address, err := common.NewAddressWitnessScriptHash(scriptHash, netParams)
-       if err != nil {
-               return ""
-       }
+// BuildP2SHAddress create P2SH address
+func BuildP2SHAddress(scriptHash []byte, netParams interface{}) string {
+       switch netParams.(type) {
+       case *consensus.Params:
+               params := netParams.(*consensus.Params)
+               address, err := common.NewAddressWitnessScriptHash(scriptHash, params)
+               if err == nil {
+                       return address.EncodeAddress()
+               }
 
-       return address.EncodeAddress()
+       case *btmConsensus.Params:
+               params := netParams.(*btmConsensus.Params)
+               address, err := btmCommon.NewAddressWitnessScriptHash(scriptHash, params)
+               if err == nil {
+                       return address.EncodeAddress()
+               }
+       }
+       return ""
 }
 
 // BuildAnnotatedOutput build the annotated output.