From 4362341b437f1699ec16333c1c19a4709cf265e0 Mon Sep 17 00:00:00 2001
From: Alexsander Falcucci <alex.falcucci@gmail.com>
Date: Sun, 11 Feb 2024 16:51:35 +0100
Subject: [PATCH] fix: support multiple pools in stake snapshot query (#396)

---
 examples/n2c-miniprotocols/src/main.rs                     | 7 +++++--
 .../src/miniprotocols/localstate/queries_v16/codec.rs      | 2 +-
 .../src/miniprotocols/localstate/queries_v16/mod.rs        | 5 +++--
 pallas-network/tests/protocols.rs                          | 5 +++--
 4 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/examples/n2c-miniprotocols/src/main.rs b/examples/n2c-miniprotocols/src/main.rs
index 60b6342..16c2483 100644
--- a/examples/n2c-miniprotocols/src/main.rs
+++ b/examples/n2c-miniprotocols/src/main.rs
@@ -1,4 +1,7 @@
+use std::collections::BTreeSet;
+
 use pallas::{
+    codec::utils::Bytes,
     ledger::{addresses::Address, traverse::MultiEraBlock},
     network::{
         facades::NodeClient,
@@ -55,8 +58,8 @@ async fn do_localstate_query(client: &mut NodeClient) {
     println!("result: {:?}", result);
 
     // Stake pool ID/verification key hash (either Bech32-decoded or hex-decoded).
-    // Empty Vec means all pools.
-    let pools = vec![];
+    // Empty Set means all pools.
+    let pools: BTreeSet<Bytes> = BTreeSet::new();
     let result = queries_v16::get_stake_snapshots(client, era, pools)
         .await
         .unwrap();
diff --git a/pallas-network/src/miniprotocols/localstate/queries_v16/codec.rs b/pallas-network/src/miniprotocols/localstate/queries_v16/codec.rs
index b8eb478..611a10f 100644
--- a/pallas-network/src/miniprotocols/localstate/queries_v16/codec.rs
+++ b/pallas-network/src/miniprotocols/localstate/queries_v16/codec.rs
@@ -99,7 +99,7 @@ impl Encode<()> for BlockQuery {
                 e.u16(20)?;
 
                 if !pools.is_empty() {
-                    e.array(Vec::len(pools) as u64)?;
+                    e.array(1)?;
                     e.tag(Tag::Unassigned(258))?;
                 }
 
diff --git a/pallas-network/src/miniprotocols/localstate/queries_v16/mod.rs b/pallas-network/src/miniprotocols/localstate/queries_v16/mod.rs
index 8165e27..0d01e81 100644
--- a/pallas-network/src/miniprotocols/localstate/queries_v16/mod.rs
+++ b/pallas-network/src/miniprotocols/localstate/queries_v16/mod.rs
@@ -1,6 +1,7 @@
 // TODO: this should move to pallas::ledger crate at some point
 
 use pallas_crypto::hash::Hash;
+use std::collections::BTreeSet;
 use std::hash::Hash as StdHash;
 // required for derive attrs to work
 use pallas_codec::minicbor::{self};
@@ -210,7 +211,7 @@ pub type Addr = Bytes;
 
 pub type Addrs = Vec<Addr>;
 
-pub type Pools = Vec<Option<Bytes>>;
+pub type Pools = BTreeSet<Bytes>;
 
 pub type Coin = AnyUInt;
 
@@ -388,7 +389,7 @@ pub async fn get_utxo_by_address(
 pub async fn get_stake_snapshots(
     client: &mut Client,
     era: u16,
-    pools: Vec<Option<Bytes>>,
+    pools: BTreeSet<Bytes>,
 ) -> Result<StakeSnapshot, ClientError> {
     let query = BlockQuery::GetStakeSnapshots(pools);
     let query = LedgerQuery::BlockQuery(era, query);
diff --git a/pallas-network/tests/protocols.rs b/pallas-network/tests/protocols.rs
index 70f8e55..0f6a67d 100644
--- a/pallas-network/tests/protocols.rs
+++ b/pallas-network/tests/protocols.rs
@@ -1,3 +1,4 @@
+use std::collections::BTreeSet;
 use std::fs;
 use std::net::{Ipv4Addr, SocketAddrV4};
 use std::time::Duration;
@@ -673,7 +674,7 @@ pub async fn local_state_query_server_and_client_happy_path() {
                 localstate::queries_v16::Request::LedgerQuery(
                     localstate::queries_v16::LedgerQuery::BlockQuery(
                         5,
-                        localstate::queries_v16::BlockQuery::GetStakeSnapshots(vec![]),
+                        localstate::queries_v16::BlockQuery::GetStakeSnapshots(BTreeSet::new()),
                     ),
                 )
             );
@@ -923,7 +924,7 @@ pub async fn local_state_query_server_and_client_happy_path() {
         let request = AnyCbor::from_encode(localstate::queries_v16::Request::LedgerQuery(
             localstate::queries_v16::LedgerQuery::BlockQuery(
                 5,
-                localstate::queries_v16::BlockQuery::GetStakeSnapshots(vec![]),
+                localstate::queries_v16::BlockQuery::GetStakeSnapshots(BTreeSet::new()),
             ),
         ));