From fe34c88dfcaf32c2ce91df9b7cb34935aa9f3b8d Mon Sep 17 00:00:00 2001
From: Ivan Kuznetsov <me@jeiwan.ru>
Date: Sun, 17 Sep 2017 12:19:01 +0700
Subject: [PATCH] Implement UTXOSet.Update

---
 utxo_set.go | 184 +++++++++++++++++++++++++++++++++-------------------
 1 file changed, 119 insertions(+), 65 deletions(-)

diff --git a/utxo_set.go b/utxo_set.go
index 74e7fec..c85cf92 100644
--- a/utxo_set.go
+++ b/utxo_set.go
@@ -14,6 +14,87 @@ type UTXOSet struct {
 	Blockchain *Blockchain
 }
 
+// FindSpendableOutputs finds and returns unspent outputs to reference in inputs
+func (u UTXOSet) FindSpendableOutputs(pubkeyHash []byte, amount int) (int, map[string][]int) {
+	unspentOutputs := make(map[string][]int)
+	accumulated := 0
+	db := u.Blockchain.db
+
+	err := db.View(func(tx *bolt.Tx) error {
+		b := tx.Bucket([]byte(utxoBucket))
+		c := b.Cursor()
+
+		for k, v := c.First(); k != nil; k, v = c.Next() {
+			txID := hex.EncodeToString(k)
+			outs := DeserializeOutputs(v)
+
+			for outIdx, out := range outs.Outputs {
+				if out.IsLockedWithKey(pubkeyHash) && accumulated < amount {
+					accumulated += out.Value
+					unspentOutputs[txID] = append(unspentOutputs[txID], outIdx)
+				}
+			}
+		}
+
+		return nil
+	})
+	if err != nil {
+		log.Panic(err)
+	}
+
+	return accumulated, unspentOutputs
+}
+
+// FindUTXO finds UTXO for a public key hash
+func (u UTXOSet) FindUTXO(pubKeyHash []byte) []TXOutput {
+	var UTXOs []TXOutput
+	db := u.Blockchain.db
+
+	err := db.View(func(tx *bolt.Tx) error {
+		b := tx.Bucket([]byte(utxoBucket))
+		c := b.Cursor()
+
+		for k, v := c.First(); k != nil; k, v = c.Next() {
+			outs := DeserializeOutputs(v)
+
+			for _, out := range outs.Outputs {
+				if out.IsLockedWithKey(pubKeyHash) {
+					UTXOs = append(UTXOs, out)
+				}
+			}
+		}
+
+		return nil
+	})
+	if err != nil {
+		log.Panic(err)
+	}
+
+	return UTXOs
+}
+
+// GetCount returns the number of transactions in the UTXO set
+func (u UTXOSet) GetCount() int {
+	db := u.Blockchain.db
+	counter := 0
+
+	err := db.View(func(tx *bolt.Tx) error {
+		b := tx.Bucket([]byte(utxoBucket))
+		c := b.Cursor()
+
+		for k, _ := c.First(); k != nil; k, _ = c.Next() {
+			counter++
+		}
+
+		return nil
+	})
+	if err != nil {
+		log.Panic(err)
+	}
+
+	return counter
+}
+
 // Reindex rebuilds the UTXO set
 func (u UTXOSet) Reindex() {
 	db := u.Blockchain.db
@@ -61,45 +142,51 @@ func (u UTXOSet) Reindex() {
 	})
 }
 
-// GetCount returns the number of transactions in the UTXO set
-func (u UTXOSet) GetCount() int {
-	db := u.Blockchain.db
-	counter := 0
-
-	err := db.View(func(tx *bolt.Tx) error {
-		b := tx.Bucket([]byte(utxoBucket))
-		c := b.Cursor()
-
-		for k, _ := c.First(); k != nil; k, _ = c.Next() {
-			counter++
-		}
-
-		return nil
-	})
-	if err != nil {
-		log.Panic(err)
-	}
-
-	return counter
-}
-
-// FindUTXO finds UTXO for a public key hash
-func (u UTXOSet) FindUTXO(pubKeyHash []byte) []TXOutput {
-	var UTXOs []TXOutput
+// Update updates the UTXO set with transactions from the Block
+// The Block is considered to be the tip of a blockchain
+func (u UTXOSet) Update(block *Block) {
 	db := u.Blockchain.db
 
-	err := db.View(func(tx *bolt.Tx) error {
+	err := db.Update(func(tx *bolt.Tx) error {
 		b := tx.Bucket([]byte(utxoBucket))
-		c := b.Cursor()
 
-		for k, v := c.First(); k != nil; k, v = c.Next() {
-			outs := DeserializeOutputs(v)
+		for _, tx := range block.Transactions {
+			if tx.IsCoinbase() == false {
+				for _, vin := range tx.Vin {
+					updatedOuts := TXOutputs{}
+					data := b.Get(vin.Txid)
+					outs := DeserializeOutputs(data)
+
+					for outIdx, out := range outs.Outputs {
+						if outIdx != vin.Vout {
+							updatedOuts.Outputs = append(updatedOuts.Outputs, out)
+						}
+					}
+
+					if len(updatedOuts.Outputs) == 0 {
+						err := b.Delete(vin.Txid)
+						if err != nil {
+							log.Panic(err)
+						}
+					} else {
+						err := b.Put(vin.Txid, updatedOuts.Serialize())
+						if err != nil {
+							log.Panic(err)
+						}
+					}
 
-			for _, out := range outs.Outputs {
-				if out.IsLockedWithKey(pubKeyHash) {
-					UTXOs = append(UTXOs, out)
 				}
 			}
+
+			newOutputs := TXOutputs{}
+			for _, out := range tx.Vout {
+				newOutputs.Outputs = append(newOutputs.Outputs, out)
+			}
+
+			err := b.Put(tx.ID, newOutputs.Serialize())
+			if err != nil {
+				log.Panic(err)
+			}
 		}
 
 		return nil
@@ -107,37 +194,4 @@ func (u UTXOSet) FindUTXO(pubKeyHash []byte) []TXOutput {
 	if err != nil {
 		log.Panic(err)
 	}
-
-	return UTXOs
-}
-
-// FindSpendableOutputs finds and returns unspent outputs to reference in inputs
-func (u UTXOSet) FindSpendableOutputs(pubkeyHash []byte, amount int) (int, map[string][]int) {
-	unspentOutputs := make(map[string][]int)
-	accumulated := 0
-	db := u.Blockchain.db
-
-	err := db.View(func(tx *bolt.Tx) error {
-		b := tx.Bucket([]byte(utxoBucket))
-		c := b.Cursor()
-
-		for k, v := c.First(); k != nil; k, v = c.Next() {
-			txID := hex.EncodeToString(k)
-			outs := DeserializeOutputs(v)
-
-			for outIdx, out := range outs.Outputs {
-				if out.IsLockedWithKey(pubkeyHash) && accumulated < amount {
-					accumulated += out.Value
-					unspentOutputs[txID] = append(unspentOutputs[txID], outIdx)
-				}
-			}
-		}
-
-		return nil
-	})
-	if err != nil {
-		log.Panic(err)
-	}
-
-	return accumulated, unspentOutputs
 }