Implement additional helpers in ic-nervous-system-agent

rs/nervous_system/agent/BUILD.bazel

@@ -6,10 +6,13 @@ DEPENDENCIES = [
     # Keep sorted.
     "//packages/pocket-ic",
     "//rs/nervous_system/clients",
+    "//rs/nns/common",
     "//rs/nns/constants",
+    "//rs/nns/governance/api",
     "//rs/nns/sns-wasm",
     "//rs/sns/governance",
     "//rs/sns/root",
+    "//rs/sns/swap",
     "//rs/types/base_types",
     "@crate_index//:anyhow",
     "@crate_index//:candid",

rs/nervous_system/agent/Cargo.toml

@@ -13,11 +13,14 @@ candid = { workspace = true }
 ic-agent = { workspace = true }
 ic-base-types = { path = "../../types/base_types" }
 ic-nervous-system-clients = { path = "../clients" }
+ic-nns-governance-api = { path = "../../nns/governance/api" }
+ic-nns-common = { path = "../../nns/common" }
 ic-nns-constants = { path = "../../nns/constants" }
 ic-sns-wasm = { path = "../../nns/sns-wasm" }
 ic-sns-governance = { path = "../../sns/governance" }
 pocket-ic = { path = "../../../packages/pocket-ic" }
 ic-sns-root = { path = "../../sns/root" }
+ic-sns-swap = { path = "../../sns/swap" }
 serde = { workspace = true }
 tempfile = { workspace = true }
 thiserror = { workspace = true }

rs/nervous_system/agent/src/agent_impl.rs

@@ -15,6 +15,8 @@ pub enum AgentCallError {
     CandidDecode(candid::Error),
 }
 
+impl crate::sealed::Sealed for Agent {}
+
 impl CallCanisters for Agent {
     type Error = AgentCallError;
     async fn call<R: Request>(

rs/nervous_system/agent/src/lib.rs

@@ -7,8 +7,15 @@ use candid::Principal;
 use ic_nervous_system_clients::Request;
 use std::fmt::Display;
 
-pub trait CallCanisters {
-    type Error: Display + Send;
+// This is used to "seal" the CallCanisters trait so that it cannot be implemented outside of this crate.
+// This is useful because it means we can modify the trait in the future without worrying about
+// breaking backwards compatibility with implementations outside of this crate.
+mod sealed {
+    pub trait Sealed {}
+}
+
+pub trait CallCanisters: sealed::Sealed {
+    type Error: Display + Send + std::error::Error + 'static;
     fn call<R: Request>(
         &self,
         canister_id: impl Into<Principal> + Send,

rs/nervous_system/agent/src/nns/governance.rs

@@ -0,0 +1,16 @@
+use crate::CallCanisters;
+use ic_nns_common::pb::v1::ProposalId;
+use ic_nns_constants::GOVERNANCE_CANISTER_ID;
+use ic_nns_governance_api::pb::v1::{
+    GetNeuronsFundAuditInfoRequest, GetNeuronsFundAuditInfoResponse,
+};
+
+pub async fn get_neurons_fund_audit_info<C: CallCanisters>(
+    agent: &C,
+    nns_proposal_id: ProposalId,
+) -> Result<GetNeuronsFundAuditInfoResponse, C::Error> {
+    let request = GetNeuronsFundAuditInfoRequest {
+        nns_proposal_id: Some(nns_proposal_id),
+    };
+    agent.call(GOVERNANCE_CANISTER_ID, request).await
+}

rs/nervous_system/agent/src/nns/mod.rs

@@ -1 +1,2 @@
+pub mod governance;
 pub mod sns_wasm;

rs/nervous_system/agent/src/pocketic_impl.rs

@@ -17,6 +17,8 @@ pub enum PocketIcCallError {
     CandidDecode(candid::Error),
 }
 
+impl crate::sealed::Sealed for PocketIc {}
+
 impl CallCanisters for PocketIc {
     type Error = PocketIcCallError;
     async fn call<R: Request>(

rs/nervous_system/agent/src/sns/governance.rs

@@ -28,3 +28,10 @@ impl GovernanceCanister {
             .await
     }
 }
+
+impl GovernanceCanister {
+    pub fn new(canister_id: impl Into<PrincipalId>) -> Self {
+        let canister_id = canister_id.into();
+        Self { canister_id }
+    }
+}

rs/nervous_system/agent/src/sns/index.rs

@@ -5,3 +5,10 @@ use serde::{Deserialize, Serialize};
 pub struct IndexCanister {
     pub canister_id: PrincipalId,
 }
+
+impl IndexCanister {
+    pub fn new(canister_id: impl Into<PrincipalId>) -> Self {
+        let canister_id = canister_id.into();
+        Self { canister_id }
+    }
+}

rs/nervous_system/agent/src/sns/ledger.rs

@@ -5,3 +5,10 @@ use serde::{Deserialize, Serialize};
 pub struct LedgerCanister {
     pub canister_id: PrincipalId,
 }
+
+impl LedgerCanister {
+    pub fn new(canister_id: impl Into<PrincipalId>) -> Self {
+        let canister_id = canister_id.into();
+        Self { canister_id }
+    }
+}

rs/nervous_system/agent/src/sns/root.rs

@@ -10,6 +10,11 @@ pub struct RootCanister {
 }
 
 impl RootCanister {
+    pub fn new(canister_id: impl Into<PrincipalId>) -> Self {
+        let canister_id = canister_id.into();
+        Self { canister_id }
+    }
+
     pub async fn sns_canisters_summary<C: CallCanisters>(
         &self,
         agent: &C,

rs/nervous_system/agent/src/sns/swap.rs

@@ -1,7 +1,82 @@
 use ic_base_types::PrincipalId;
+use ic_sns_swap::pb::v1::{
+    GetDerivedStateRequest, GetDerivedStateResponse, GetInitRequest, GetInitResponse,
+    ListSnsNeuronRecipesRequest, ListSnsNeuronRecipesResponse, SnsNeuronRecipe,
+};
 use serde::{Deserialize, Serialize};
+use thiserror::Error;
+
+use crate::CallCanisters;
 
 #[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct SwapCanister {
     pub canister_id: PrincipalId,
 }
+
+#[derive(Debug, Error)]
+pub enum ListAllSnsNeuronRecipesError<CallCanistersError: std::error::Error + 'static> {
+    #[error(transparent)]
+    CanisterCallError(#[from] CallCanistersError),
+    #[error("There seem to be too many neuron recipes ({0}).")]
+    TooManyRecipes(u64),
+}
+
+impl SwapCanister {
+    pub fn new(canister_id: impl Into<PrincipalId>) -> Self {
+        let canister_id = canister_id.into();
+        Self { canister_id }
+    }
+
+    pub async fn get_derived_state<C: CallCanisters>(
+        &self,
+        agent: &C,
+    ) -> Result<GetDerivedStateResponse, C::Error> {
+        agent
+            .call(self.canister_id, GetDerivedStateRequest {})
+            .await
+    }
+
+    pub async fn get_init<C: CallCanisters>(&self, agent: &C) -> Result<GetInitResponse, C::Error> {
+        agent.call(self.canister_id, GetInitRequest {}).await
+    }
+
+    pub async fn list_sns_neuron_recipes<C: CallCanisters>(
+        &self,
+        agent: &C,
+        limit: u32,
+        offset: u64,
+    ) -> Result<ListSnsNeuronRecipesResponse, C::Error> {
+        agent
+            .call(
+                self.canister_id,
+                ListSnsNeuronRecipesRequest {
+                    limit: Some(limit),
+                    offset: Some(offset),
+                },
+            )
+            .await
+    }
+
+    pub async fn list_all_sns_neuron_recipes<C: CallCanisters>(
+        &self,
+        agent: &C,
+    ) -> Result<Vec<SnsNeuronRecipe>, ListAllSnsNeuronRecipesError<C::Error>> {
+        let mut sns_neuron_recipes: Vec<SnsNeuronRecipe> = vec![];
+        let batch_size = 10_000_u64;
+        let num_calls = 100_u64;
+        for i in 0..num_calls {
+            let new_sns_neuron_recipes = self
+                .list_sns_neuron_recipes(agent, batch_size as u32, batch_size * i)
+                .await
+                .map_err(ListAllSnsNeuronRecipesError::CanisterCallError)?;
+            if new_sns_neuron_recipes.sns_neuron_recipes.is_empty() {
+                return Ok(sns_neuron_recipes);
+            } else {
+                sns_neuron_recipes.extend(new_sns_neuron_recipes.sns_neuron_recipes.into_iter())
+            }
+        }
+        Err(ListAllSnsNeuronRecipesError::TooManyRecipes(
+            batch_size * num_calls,
+        ))
+    }
+}

rs/nns/governance/api/src/lib.rs

@@ -2,5 +2,6 @@ pub mod bitcoin;
 pub mod pb;
 pub mod proposal_submission_helpers;
 pub mod proposal_validation;
+mod request_impls;
 pub mod subnet_rental;
 pub mod test_api;

rs/nns/governance/api/src/request_impls.rs

@@ -0,0 +1,7 @@
+use ic_nervous_system_clients::Request;
+
+impl Request for crate::pb::v1::GetNeuronsFundAuditInfoRequest {
+    type Response = crate::pb::v1::GetNeuronsFundAuditInfoResponse;
+    const METHOD: &'static str = "get_neurons_fund_audit_info";
+    const UPDATE: bool = false;
+}

rs/sns/swap/src/lib.rs

@@ -4,6 +4,7 @@ pub mod logs;
 pub mod memory;
 pub mod neurons_fund;
 pub mod pb;
+mod request_impls;
 pub mod swap;
 pub mod swap_builder;
 pub mod types;

rs/sns/swap/src/request_impls.rs

@@ -0,0 +1,19 @@
+use ic_nervous_system_clients::Request;
+
+impl Request for crate::pb::v1::GetDerivedStateRequest {
+    type Response = crate::pb::v1::GetDerivedStateResponse;
+    const METHOD: &'static str = "get_derived_state";
+    const UPDATE: bool = false;
+}
+
+impl Request for crate::pb::v1::GetInitRequest {
+    type Response = crate::pb::v1::GetInitResponse;
+    const METHOD: &'static str = "get_init";
+    const UPDATE: bool = false;
+}
+
+impl Request for crate::pb::v1::ListSnsNeuronRecipesRequest {
+    type Response = crate::pb::v1::ListSnsNeuronRecipesResponse;
+    const METHOD: &'static str = "list_sns_neuron_recipes";
+    const UPDATE: bool = false;
+}