From 010c32595d9f083f8c25d0e5f2c2e3a9b053047c Mon Sep 17 00:00:00 2001
From: Aurora Gaffney <aurora@blinklabs.io>
Date: Sun, 25 Aug 2024 19:41:45 -0500
Subject: [PATCH] feat: ledger functions to query UTxOs (#99)

---
 state/models/utxo.go | 29 +++++++++++++++++++++++++++++
 state/state.go       | 21 +++++++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/state/models/utxo.go b/state/models/utxo.go
index 43879b4..3bec856 100644
--- a/state/models/utxo.go
+++ b/state/models/utxo.go
@@ -21,6 +21,7 @@ import (
 	"github.com/blinklabs-io/gouroboros/ledger"
 	"github.com/blinklabs-io/node/database"
 	"github.com/dgraph-io/badger/v4"
+	"gorm.io/gorm"
 )
 
 type Utxo struct {
@@ -69,6 +70,34 @@ func UtxoByRef(db database.Database, txId []byte, outputIdx uint32) (Utxo, error
 	return tmpUtxo, result.Error
 }
 
+func UtxosByAddress(db database.Database, addr ledger.Address) ([]Utxo, error) {
+	var ret []Utxo
+	// Build sub-query for address
+	var addrQuery *gorm.DB
+	if addr.PaymentKeyHash() != ledger.NewBlake2b224(nil) {
+		addrQuery = db.Metadata().Where("payment_key = ?", addr.PaymentKeyHash().Bytes())
+	}
+	if addr.StakeKeyHash() != ledger.NewBlake2b224(nil) {
+		if addrQuery != nil {
+			addrQuery = addrQuery.Or("staking_key = ?", addr.StakeKeyHash().Bytes())
+		} else {
+			addrQuery = db.Metadata().Where("staking_key = ?", addr.StakeKeyHash().Bytes())
+		}
+	}
+	result := db.Metadata().Where("deleted_slot = 0").Where(addrQuery).Find(&ret)
+	if result.Error != nil {
+		return nil, result.Error
+	}
+	// Load CBOR from blob DB for each UTxO
+	for idx, tmpUtxo := range ret {
+		if err := tmpUtxo.loadCbor(db.Blob()); err != nil {
+			return nil, err
+		}
+		ret[idx] = tmpUtxo
+	}
+	return ret, nil
+}
+
 func UtxoDelete(db database.Database, utxo Utxo) error {
 	// Remove from metadata DB
 	if result := db.Metadata().Delete(&utxo); result.Error != nil {
diff --git a/state/state.go b/state/state.go
index c830b5e..e424acb 100644
--- a/state/state.go
+++ b/state/state.go
@@ -276,6 +276,10 @@ func (ls *LedgerState) consumeUtxo(utxoId ledger.TransactionInput, slot uint64)
 	// Find UTxO
 	utxo, err := models.UtxoByRef(ls.db, utxoId.Id().Bytes(), utxoId.Index())
 	if err != nil {
+		// TODO: make this configurable?
+		if err == gorm.ErrRecordNotFound {
+			return nil
+		}
 		return err
 	}
 	// Mark as deleted in specified slot
@@ -348,6 +352,7 @@ func (ls *LedgerState) RecentChainPoints(count int) ([]ocommon.Point, error) {
 	return ret, nil
 }
 
+// GetIntersectPoint returns the intersect between the specified points and the current chain
 func (ls *LedgerState) GetIntersectPoint(points []ocommon.Point) (*ocommon.Point, error) {
 	var ret ocommon.Point
 	for _, point := range points {
@@ -373,10 +378,12 @@ func (ls *LedgerState) GetIntersectPoint(points []ocommon.Point) (*ocommon.Point
 	return nil, nil
 }
 
+// GetChainFromPoint returns a ChainIterator starting at the specified point
 func (ls *LedgerState) GetChainFromPoint(point ocommon.Point) (*ChainIterator, error) {
 	return newChainIterator(ls, point)
 }
 
+// Tip returns the current chain tip
 func (ls *LedgerState) Tip() (ochainsync.Tip, error) {
 	var ret ochainsync.Tip
 	var tmpBlock models.Block
@@ -390,3 +397,17 @@ func (ls *LedgerState) Tip() (ochainsync.Tip, error) {
 	}
 	return ret, nil
 }
+
+// UtxoByRef returns a single UTxO by reference
+func (ls *LedgerState) UtxoByRef(txId []byte, outputIdx uint32) (models.Utxo, error) {
+	ls.RLock()
+	defer ls.RUnlock()
+	return models.UtxoByRef(ls.db, txId, outputIdx)
+}
+
+// UtxosByAddress returns all UTxOs that belong to the specified address
+func (ls *LedgerState) UtxosByAddress(addr ledger.Address) ([]models.Utxo, error) {
+	ls.RLock()
+	defer ls.RUnlock()
+	return models.UtxosByAddress(ls.db, addr)
+}