Fix per factor source

radixdlt · Dec 4, 2024 · abb7972 · abb7972
1 parent 4c6f67b
commit abb7972
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 8 deletions.
diff --git a/crates/sargon-uniffi/src/keys_collector/key_derivation_request.rs b/crates/sargon-uniffi/src/keys_collector/key_derivation_request.rs
@@ -2,6 +2,7 @@ use crate::prelude::*;
 use sargon::indexmap::IndexMap;
 use sargon::{IndexSet, KeyDerivationRequest as InternalKeyDerivationRequest};
 
+/// A collection of derivation paths, on a per-factor-source basis.
 #[derive(Clone, PartialEq, Eq, uniffi::Record)]
 pub struct KeyDerivationRequest {
     pub per_factor_source: HashMap<FactorSourceIDFromHash, Vec<DerivationPath>>,

diff --git a/crates/sargon-uniffi/src/keys_collector/key_derivation_response.rs b/crates/sargon-uniffi/src/keys_collector/key_derivation_response.rs
@@ -3,14 +3,30 @@ use sargon::IndexMap;
 use sargon::IndexSet;
 use sargon::KeyDerivationResponse as InternalKeyDerivationResponse;
 
+/// A collection of `HierarchicalDeterministicFactorInstance`s, on a
+/// per-factor-source basis. In case of MonoKeyDerivation the list will contain
+/// a single `KeyDerivationPerFactorSource`.
 #[derive(Clone, PartialEq, Eq, uniffi::Record)]
 pub struct KeyDerivationResponse {
-    pub per_factor_source: HashMap<
-        FactorSourceIDFromHash,
-        Vec<HierarchicalDeterministicFactorInstance>,
-    >,
+    pub per_factor_source: Vec<KeyDerivationPerFactorSource>,
 }
 
+#[derive(Clone, PartialEq, Eq, uniffi::Record)]
+pub struct KeyDerivationPerFactorSource {
+    pub factor_source_id: FactorSourceIDFromHash,
+    pub factor_instances: Vec<HierarchicalDeterministicFactorInstance>,
+}
+
+impl KeyDerivationPerFactorSource {
+    pub fn new(
+        factor_source_id: FactorSourceIDFromHash,
+        factor_instances: Vec<HierarchicalDeterministicFactorInstance>,
+    ) -> Self {
+        Self { factor_source_id, factor_instances }
+    }
+}
+
+
 impl KeyDerivationResponse {
     pub fn into_internal(&self) -> InternalKeyDerivationResponse {
         self.clone().into()
@@ -24,7 +40,10 @@ impl From<InternalKeyDerivationResponse> for KeyDerivationResponse {
                 .per_factor_source
                 .into_iter()
                 .map(|(k, v)| {
-                    (k.into(), v.into_iter().map(|d| d.into()).collect())
+                    KeyDerivationPerFactorSource::new(
+                        k.into(),
+                        v.into_iter().map(|d| d.into()).collect()
+                    )
                 })
                 .collect(),
         }
@@ -34,11 +53,11 @@ impl From<InternalKeyDerivationResponse> for KeyDerivationResponse {
 impl From<KeyDerivationResponse> for InternalKeyDerivationResponse {
     fn from(value: KeyDerivationResponse) -> Self {
         Self::new(IndexMap::from_iter(
-            value.per_factor_source.into_iter().map(|(k, v)| {
+            value.per_factor_source.into_iter().map(|item| {
                 (
-                    k.into_internal(),
+                    item.factor_source_id.into_internal(),
                     IndexSet::from_iter(
-                        v.into_iter().map(|d| d.into_internal()),
+                        item.factor_instances.into_iter().map(|d| d.into_internal()),
                     ),
                 )
             }),