From 4c40f9f514d520d7b3692c79951458092fcae809 Mon Sep 17 00:00:00 2001
From: Alfredo Garcia <oxarbitrage@gmail.com>
Date: Sun, 19 Apr 2020 20:22:06 -0300
Subject: [PATCH] implement z_getnotescount api call

---
 qa/rpc-tests/wallet_listreceived.py |  8 ++++++
 src/rpc/client.cpp                  |  1 +
 src/wallet/rpcwallet.cpp            | 44 +++++++++++++++++++++++++++++
 3 files changed, 53 insertions(+)

diff --git a/qa/rpc-tests/wallet_listreceived.py b/qa/rpc-tests/wallet_listreceived.py
index 07004a4ff..52dff938a 100755
--- a/qa/rpc-tests/wallet_listreceived.py
+++ b/qa/rpc-tests/wallet_listreceived.py
@@ -98,6 +98,8 @@ class ListReceivedTest (BitcoinTestFramework):
 
         r = self.nodes[1].z_listreceivedbyaddress(zaddr1)
         assert_equal(0, len(r), "Should have received no confirmed note")
+        c = self.nodes[1].z_getnotescount()
+        assert_equal(0, c[release], "Count of confirmed notes should be 0")
 
         # No confirmation required, one note should be present
         r = self.nodes[1].z_listreceivedbyaddress(zaddr1, 0)
@@ -110,6 +112,9 @@ class ListReceivedTest (BitcoinTestFramework):
         assert_equal(-1, r[0]['blockindex'])
         assert_equal(0, r[0]['blockheight'])
 
+        c = self.nodes[1].z_getnotescount(0)
+        assert_equal(1, c[release], "Count of unconfirmed notes should be 1")
+
         # Confirm transaction (1 ZEC from taddr to zaddr1)
         self.generate_and_sync(height+3)
 
@@ -215,6 +220,9 @@ class ListReceivedTest (BitcoinTestFramework):
         assert_false(r[0]['change'], "Note valued at 0.6 should not be change")
         assert_equal(no_memo, r[0]['memo'])
 
+        c = self.nodes[1].z_getnotescount(0)
+        assert_equal(3, c[release], "Count of unconfirmed notes should be 3(2 in zaddr1 + 1 in zaddr2)")
+
     def run_test(self):
         self.run_test_release('sprout', 200)
         self.run_test_release('sapling', 214)
diff --git a/src/rpc/client.cpp b/src/rpc/client.cpp
index 950e3269c..fe084b8f2 100644
--- a/src/rpc/client.cpp
+++ b/src/rpc/client.cpp
@@ -144,6 +144,7 @@ static const CRPCConvertParam vRPCConvertParams[] =
     { "z_getpaymentdisclosure", 1},
     { "z_getpaymentdisclosure", 2},
     { "z_setmigration", 0},
+    { "z_getnotescount", 0},
 };
 
 class CRPCConvertTable
diff --git a/src/wallet/rpcwallet.cpp b/src/wallet/rpcwallet.cpp
index 508a9282a..e4b88d2c5 100644
--- a/src/wallet/rpcwallet.cpp
+++ b/src/wallet/rpcwallet.cpp
@@ -4967,6 +4967,49 @@ UniValue z_listoperationids(const UniValue& params, bool fHelp)
     return ret;
 }
 
+
+UniValue z_getnotescount(const UniValue& params, bool fHelp)
+{
+    if (!EnsureWalletIsAvailable(fHelp))
+        return NullUniValue;
+
+    if (fHelp || params.size() > 1)
+        throw runtime_error(
+            "z_getnotescount\n"
+            "\nArguments:\n"
+            "1. minconf      (numeric, optional, default=1) Only include notes in transactions confirmed at least this many times.\n"
+            "\nReturns the number of sprout and sapling notes available in the wallet.\n"
+            "\nResult:\n"
+            "{\n"
+            "  \"sprout\"      (numeric) the number of sprout notes in the wallet\n"
+            "  \"sapling\"     (numeric) the number of sapling notes in the wallet\n"
+            "}\n"
+            "\nExamples:\n"
+            + HelpExampleCli("z_getnotescount", "0")
+            + HelpExampleRpc("z_getnotescount", "0")
+        );
+
+    LOCK2(cs_main, pwalletMain->cs_wallet);
+
+    int nMinDepth = 1;
+    if (params.size() > 0)
+        nMinDepth = params[0].get_int();
+
+    int sprout = 0;
+    int sapling = 0;
+    for (auto& wtx : pwalletMain->mapWallet) {
+        if (wtx.second.GetDepthInMainChain() >= nMinDepth) {
+            sprout += wtx.second.mapSproutNoteData.size();
+            sapling += wtx.second.mapSaplingNoteData.size();
+        }
+    }
+    UniValue ret(UniValue::VOBJ);
+    ret.push_back(Pair("sprout", sprout));
+    ret.push_back(Pair("sapling", sapling));
+
+    return ret;
+}
+
 extern UniValue dumpprivkey(const UniValue& params, bool fHelp); // in rpcdump.cpp
 extern UniValue importprivkey(const UniValue& params, bool fHelp);
 extern UniValue importaddress(const UniValue& params, bool fHelp);
@@ -5054,6 +5097,7 @@ static const CRPCCommand commands[] =
     { "wallet",             "z_exportwallet",           &z_exportwallet,           true  },
     { "wallet",             "z_importwallet",           &z_importwallet,           true  },
     { "wallet",             "z_viewtransaction",        &z_viewtransaction,        false },
+    { "wallet",             "z_getnotescount",          &z_getnotescount,          false },
     // TODO: rearrange into another category
     { "disclosure",         "z_getpaymentdisclosure",   &z_getpaymentdisclosure,   true  },
     { "disclosure",         "z_validatepaymentdisclosure", &z_validatepaymentdisclosure, true }