Spend 1 key vault

crates/chia-sdk-driver/src/driver_error.rs

@@ -53,4 +53,10 @@ pub enum DriverError {
 
     #[error("unknown puzzle")]
     UnknownPuzzle,
+
+    #[error("wrong number of spends")]
+    WrongSpendCount,
+
+    #[error("missing member spend")]
+    MissingMemberSpend,
 }

crates/chia-sdk-driver/src/primitives/vault/m_of_n.rs

@@ -1,9 +1,15 @@
+use std::collections::HashMap;
+
 use chia_protocol::Bytes32;
-use chia_sdk_types::{MerkleTree, Mod, Vault1ofNArgs, VaultMofNArgs, VaultNofNArgs};
-use clvm_utils::TreeHash;
+use chia_sdk_types::{
+    MerkleTree, Mod, Vault1ofNArgs, Vault1ofNSolution, VaultMofNArgs, VaultMofNSolution,
+    VaultNofNArgs, VaultNofNSolution,
+};
+use clvm_traits::clvm_tuple;
+use clvm_utils::{tree_hash_atom, tree_hash_pair, TreeHash};
 use clvmr::NodePtr;
 
-use crate::{DriverError, SpendContext};
+use crate::{DriverError, Spend, SpendContext};
 
 use super::{KnownPuzzles, Member, PuzzleWithRestrictions, VaultLayer};
 
@@ -35,6 +41,60 @@ impl MofN {
         &self.members
     }
 
+    pub fn solve(
+        &self,
+        ctx: &mut SpendContext,
+        member_spends: HashMap<TreeHash, Spend>,
+    ) -> Result<NodePtr, DriverError> {
+        if member_spends.len() != self.required {
+            return Err(DriverError::WrongSpendCount);
+        }
+
+        match self.optimization() {
+            Some(MofNOptimization::One) => {
+                let (member_puzzle_hash, member_spend) = member_spends
+                    .into_iter()
+                    .next()
+                    .expect("missing single spend");
+
+                let merkle_tree = self.merkle_tree();
+                let merkle_proof = merkle_tree
+                    .proof(member_puzzle_hash.into())
+                    .ok_or(DriverError::InvalidMerkleProof)?;
+
+                ctx.alloc(&Vault1ofNSolution::new(
+                    merkle_proof,
+                    member_spend.puzzle,
+                    member_spend.solution,
+                ))
+            }
+            Some(MofNOptimization::All) => {
+                let mut member_solutions = Vec::with_capacity(self.required);
+
+                for member in &self.members {
+                    let spend = member_spends
+                        .get(&member.puzzle_hash())
+                        .ok_or(DriverError::MissingMemberSpend)?;
+
+                    member_solutions.push(spend.solution);
+                }
+
+                ctx.alloc(&VaultNofNSolution::new(member_solutions))
+            }
+            None => {
+                let puzzle_hashes: Vec<Bytes32> = self
+                    .members
+                    .iter()
+                    .map(|member| member.puzzle_hash().into())
+                    .collect();
+
+                let proof = m_of_n_proof(ctx, &puzzle_hashes, &member_spends)?;
+
+                ctx.alloc(&VaultMofNSolution::new(proof))
+            }
+        }
+    }
+
     fn optimization(&self) -> Option<MofNOptimization> {
         if self.required == 1 {
             Some(MofNOptimization::One)
@@ -104,3 +164,35 @@ impl VaultLayer for MofN {
         Self { required, members }
     }
 }
+
+fn m_of_n_proof(
+    ctx: &mut SpendContext,
+    puzzle_hashes: &[Bytes32],
+    member_spends: &HashMap<TreeHash, Spend>,
+) -> Result<NodePtr, DriverError> {
+    if puzzle_hashes.len() == 1 {
+        let puzzle_hash = puzzle_hashes[0];
+
+        return if let Some(spend) = member_spends.get(&puzzle_hash.into()) {
+            ctx.alloc(&clvm_tuple!((), spend.puzzle, spend.solution))
+        } else {
+            ctx.alloc(&Bytes32::from(tree_hash_atom(&puzzle_hash)))
+        };
+    }
+
+    let mid_index = puzzle_hashes.len().div_ceil(2);
+    let first = &puzzle_hashes[..mid_index];
+    let rest = &puzzle_hashes[mid_index..];
+
+    let first_proof = m_of_n_proof(ctx, first, member_spends)?;
+    let rest_proof = m_of_n_proof(ctx, rest, member_spends)?;
+
+    if first_proof.is_pair() || rest_proof.is_pair() {
+        ctx.alloc(&(first_proof, rest_proof))
+    } else {
+        let first_hash = ctx.extract::<Bytes32>(first_proof)?;
+        let rest_hash = ctx.extract::<Bytes32>(rest_proof)?;
+        let pair_hash = Bytes32::from(tree_hash_pair(first_hash.into(), rest_hash.into()));
+        ctx.alloc(&pair_hash)
+    }
+}

crates/chia-sdk-driver/src/primitives/vault/puzzle_with_restrictions.rs

@@ -1,8 +1,11 @@
-use chia_sdk_types::{DelegatedFeederArgs, IndexWrapperArgs, Mod, RestrictionsArgs};
+use chia_sdk_types::{
+    DelegatedFeederArgs, DelegatedFeederSolution, IndexWrapperArgs, Mod, RestrictionsArgs,
+    RestrictionsSolution,
+};
 use clvm_utils::TreeHash;
 use clvmr::NodePtr;
 
-use crate::{DriverError, SpendContext};
+use crate::{DriverError, Spend, SpendContext};
 
 use super::{KnownPuzzles, Restriction, VaultLayer};
 
@@ -32,6 +35,35 @@ impl<T> PuzzleWithRestrictions<T> {
             has_delegated_feeder: false,
         }
     }
+
+    pub fn solve(
+        &self,
+        ctx: &mut SpendContext,
+        member_validator_solutions: Vec<NodePtr>,
+        delegated_puzzle_validator_solutions: Vec<NodePtr>,
+        inner_solution: NodePtr,
+        delegated_spend: Option<Spend>,
+    ) -> Result<NodePtr, DriverError> {
+        let mut solution = inner_solution;
+
+        if !self.restrictions.is_empty() {
+            solution = ctx.alloc(&RestrictionsSolution::new(
+                member_validator_solutions,
+                delegated_puzzle_validator_solutions,
+                solution,
+            ))?;
+        }
+
+        if let Some(delegated_spend) = delegated_spend {
+            solution = ctx.alloc(&DelegatedFeederSolution::new(
+                delegated_spend.puzzle,
+                delegated_spend.solution,
+                solution,
+            ))?;
+        }
+
+        Ok(solution)
+    }
 }
 
 impl<T> VaultLayer for PuzzleWithRestrictions<T>

crates/chia-sdk-driver/src/primitives/vault/vault_primitive.rs

@@ -48,10 +48,12 @@ impl VaultLayer for Vault {
 mod tests {
     use std::collections::HashMap;
 
+    use chia_puzzles::{singleton::SingletonSolution, EveProof, Proof};
     use chia_sdk_test::Simulator;
-    use chia_sdk_types::BlsMember;
+    use chia_sdk_types::{BlsMember, Conditions};
+    use clvm_traits::clvm_quote;
 
-    use crate::{Launcher, MemberKind, SpendContext, StandardLayer};
+    use crate::{Launcher, MemberKind, Spend, SpendContext, StandardLayer};
 
     use super::*;
 
@@ -72,18 +74,47 @@ mod tests {
         );
         let (mint_vault, vault) = Launcher::new(coin.coin_id(), 1).mint_vault(ctx, custody, ())?;
         p2.spend(ctx, coin, mint_vault)?;
-        sim.spend_coins(ctx.take(), &[sk])?;
+        sim.spend_coins(ctx.take(), &[sk.clone()])?;
 
         let mut known_members = HashMap::new();
         known_members.insert(
             hidden_member.curry_tree_hash(),
             MemberKind::Bls(hidden_member),
         );
-        let _vault = vault.replace(&KnownPuzzles {
+        let vault = vault.replace(&KnownPuzzles {
             members: known_members,
             ..Default::default()
         });
 
+        let delegated_puzzle = ctx.alloc(&clvm_quote!(Conditions::new().create_coin(
+            vault.custody.puzzle_hash().into(),
+            vault.coin.amount,
+            None
+        )))?;
+
+        let puzzle = vault.puzzle(ctx)?;
+        let inner_solution = vault.custody.solve(
+            ctx,
+            Vec::new(),
+            Vec::new(),
+            NodePtr::NIL,
+            Some(Spend {
+                puzzle: delegated_puzzle,
+                solution: NodePtr::NIL,
+            }),
+        )?;
+        let solution = ctx.alloc(&SingletonSolution {
+            lineage_proof: Proof::Eve(EveProof {
+                parent_parent_coin_info: coin.coin_id(),
+                parent_amount: 1,
+            }),
+            amount: 1,
+            inner_solution,
+        })?;
+
+        ctx.spend(vault.coin, Spend::new(puzzle, solution))?;
+        sim.spend_coins(ctx.take(), &[sk])?;
+
         Ok(())
     }
 }

crates/chia-sdk-types/src/puzzles/vault/restrictions.rs

@@ -2,7 +2,6 @@ mod timelock;
 
 pub use timelock::*;
 
-use chia_protocol::Bytes32;
 use clvm_traits::{FromClvm, ToClvm};
 use clvm_utils::TreeHash;
 use hex_literal::hex;
@@ -39,21 +38,18 @@ impl<MV, DV, I> Mod for RestrictionsArgs<MV, DV, I> {
 #[derive(Debug, Clone, PartialEq, Eq, ToClvm, FromClvm)]
 #[clvm(solution)]
 pub struct RestrictionsSolution<MV, DV, I> {
-    pub delegated_puzzle_hash: Bytes32,
     pub member_validator_solutions: Vec<MV>,
     pub delegated_puzzle_validator_solutions: Vec<DV>,
     pub inner_solution: I,
 }
 
 impl<MV, DV, I> RestrictionsSolution<MV, DV, I> {
     pub fn new(
-        delegated_puzzle_hash: Bytes32,
         member_validator_solutions: Vec<MV>,
         delegated_puzzle_validator_solutions: Vec<DV>,
         inner_solution: I,
     ) -> Self {
         Self {
-            delegated_puzzle_hash,
             member_validator_solutions,
             delegated_puzzle_validator_solutions,
             inner_solution,

crates/chia-sdk-types/src/puzzles/vault/vault_1_of_n.rs

@@ -25,21 +25,14 @@ impl Mod for Vault1ofNArgs {
 #[derive(Debug, Clone, PartialEq, Eq, ToClvm, FromClvm)]
 #[clvm(solution)]
 pub struct Vault1ofNSolution<P, S> {
-    pub delegated_puzzle_hash: Bytes32,
     pub merkle_proof: MerkleProof,
     pub member_puzzle: P,
     pub member_solution: S,
 }
 
 impl<P, S> Vault1ofNSolution<P, S> {
-    pub fn new(
-        delegated_puzzle_hash: Bytes32,
-        merkle_proof: MerkleProof,
-        member_puzzle: P,
-        member_solution: S,
-    ) -> Self {
+    pub fn new(merkle_proof: MerkleProof, member_puzzle: P, member_solution: S) -> Self {
         Self {
-            delegated_puzzle_hash,
             merkle_proof,
             member_puzzle,
             member_solution,

crates/chia-sdk-types/src/puzzles/vault/vault_m_of_n.rs

@@ -29,16 +29,12 @@ impl Mod for VaultMofNArgs {
 #[derive(Debug, Clone, Copy, PartialEq, Eq, ToClvm, FromClvm)]
 #[clvm(solution)]
 pub struct VaultMofNSolution<P> {
-    pub delegated_puzzle_hash: Bytes32,
     pub proofs: P,
 }
 
 impl<P> VaultMofNSolution<P> {
-    pub fn new(delegated_puzzle_hash: Bytes32, proofs: P) -> Self {
-        Self {
-            delegated_puzzle_hash,
-            proofs,
-        }
+    pub fn new(proofs: P) -> Self {
+        Self { proofs }
     }
 }
 

crates/chia-sdk-types/src/puzzles/vault/vault_n_of_n.rs

@@ -1,4 +1,3 @@
-use chia_protocol::Bytes32;
 use clvm_traits::{FromClvm, ToClvm};
 use clvm_utils::TreeHash;
 use hex_literal::hex;
@@ -25,16 +24,12 @@ impl<T> Mod for VaultNofNArgs<T> {
 #[derive(Debug, Clone, PartialEq, Eq, ToClvm, FromClvm)]
 #[clvm(solution)]
 pub struct VaultNofNSolution<T> {
-    pub delegated_puzzle_hash: Bytes32,
     pub member_solutions: Vec<T>,
 }
 
 impl<T> VaultNofNSolution<T> {
-    pub fn new(delegated_puzzle_hash: Bytes32, member_solutions: Vec<T>) -> Self {
-        Self {
-            delegated_puzzle_hash,
-            member_solutions,
-        }
+    pub fn new(member_solutions: Vec<T>) -> Self {
+        Self { member_solutions }
     }
 }