From 7eda539141fdc633f9ca1478d27922710509ef6a Mon Sep 17 00:00:00 2001
From: Ivan Kuznetsov <me@jeiwan.ru>
Date: Sun, 17 Sep 2017 10:30:30 +0700
Subject: [PATCH] Improve UTXOSet

---
 cli.go             | 13 ++++++++++++-
 cli_reindexutxo.go | 12 ++++++++++++
 utxo_set.go        | 26 ++++++++++++--------------
 3 files changed, 36 insertions(+), 15 deletions(-)
 create mode 100644 cli_reindexutxo.go

diff --git a/cli.go b/cli.go
index 863ffd1..72693ea 100644
--- a/cli.go
+++ b/cli.go
@@ -17,6 +17,7 @@ func (cli *CLI) printUsage() {
 	fmt.Println("  getbalance -address ADDRESS - Get balance of ADDRESS")
 	fmt.Println("  listaddresses - Lists all addresses from the wallet file")
 	fmt.Println("  printchain - Print all the blocks of the blockchain")
+	fmt.Println("  reindexutxo - Rebuilds the UTXO set")
 	fmt.Println("  send -from FROM -to TO -amount AMOUNT - Send AMOUNT of coins from FROM address to TO")
 }
 
@@ -35,8 +36,9 @@ func (cli *CLI) Run() {
 	createBlockchainCmd := flag.NewFlagSet("createblockchain", flag.ExitOnError)
 	createWalletCmd := flag.NewFlagSet("createwallet", flag.ExitOnError)
 	listAddressesCmd := flag.NewFlagSet("listaddresses", flag.ExitOnError)
-	sendCmd := flag.NewFlagSet("send", flag.ExitOnError)
 	printChainCmd := flag.NewFlagSet("printchain", flag.ExitOnError)
+	reindexUTXOCmd := flag.NewFlagSet("reindexutxo", flag.ExitOnError)
+	sendCmd := flag.NewFlagSet("send", flag.ExitOnError)
 
 	getBalanceAddress := getBalanceCmd.String("address", "", "The address to get balance for")
 	createBlockchainAddress := createBlockchainCmd.String("address", "", "The address to send genesis block reward to")
@@ -75,6 +77,11 @@ func (cli *CLI) Run() {
 		if err != nil {
 			log.Panic(err)
 		}
+	case "reindexutxo":
+		err := reindexUTXOCmd.Parse(os.Args[2:])
+		if err != nil {
+			log.Panic(err)
+		}
 	default:
 		cli.printUsage()
 		os.Exit(1)
@@ -108,6 +115,10 @@ func (cli *CLI) Run() {
 		cli.printChain()
 	}
 
+	if reindexUTXOCmd.Parsed() {
+		cli.reindexUTXO()
+	}
+
 	if sendCmd.Parsed() {
 		if *sendFrom == "" || *sendTo == "" || *sendAmount <= 0 {
 			sendCmd.Usage()
diff --git a/cli_reindexutxo.go b/cli_reindexutxo.go
new file mode 100644
index 0000000..e2d3717
--- /dev/null
+++ b/cli_reindexutxo.go
@@ -0,0 +1,12 @@
+package main
+
+import "fmt"
+
+func (cli *CLI) reindexUTXO() {
+	bc := NewBlockchain("")
+	UTXOSet := UTXOSet{bc}
+	UTXOSet.Reindex()
+
+	count := UTXOSet.GetCount()
+	fmt.Printf("Done! There are %d transactions in the UTXO set.", count)
+}
diff --git a/utxo_set.go b/utxo_set.go
index e43c32e..49c4c9f 100644
--- a/utxo_set.go
+++ b/utxo_set.go
@@ -10,16 +10,15 @@ import (
 const utxoBucket = "chainstate"
 
 // UTXOSet represents UTXO set
-type UTXOSet struct{}
+type UTXOSet struct {
+	Blockchain *Blockchain
+}
 
 // Reindex rebuilds the UTXO set
-func (u UTXOSet) Reindex(bc *Blockchain) {
-	db, err := bolt.Open(dbFile, 0600, nil)
-	if err != nil {
-		log.Panic(err)
-	}
+func (u UTXOSet) Reindex() {
+	db := u.Blockchain.db
 
-	err = db.Update(func(tx *bolt.Tx) error {
+	err := db.Update(func(tx *bolt.Tx) error {
 		bucketName := []byte(utxoBucket)
 		b := tx.Bucket(bucketName)
 
@@ -41,7 +40,7 @@ func (u UTXOSet) Reindex(bc *Blockchain) {
 		log.Panic(err)
 	}
 
-	UTXO := bc.FindAllUTXO()
+	UTXO := u.Blockchain.FindAllUTXO()
 
 	err = db.Update(func(tx *bolt.Tx) error {
 		b := tx.Bucket([]byte(utxoBucket))
@@ -64,14 +63,10 @@ func (u UTXOSet) Reindex(bc *Blockchain) {
 
 // GetCount returns the number of transactions in the UTXO set
 func (u UTXOSet) GetCount() int {
+	db := u.Blockchain.db
 	counter := 0
 
-	db, err := bolt.Open(dbFile, 0600, nil)
-	if err != nil {
-		log.Panic(err)
-	}
-
-	err = db.View(func(tx *bolt.Tx) error {
+	err := db.View(func(tx *bolt.Tx) error {
 		b := tx.Bucket([]byte(utxoBucket))
 		c := b.Cursor()
 
@@ -81,6 +76,9 @@ func (u UTXOSet) GetCount() int {
 
 		return nil
 	})
+	if err != nil {
+		log.Panic(err)
+	}
 
 	return counter
 }