Allocate actual puzzle

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

@@ -50,4 +50,7 @@ pub enum DriverError {
 
     #[error("invalid merkle proof")]
     InvalidMerkleProof,
+
+    #[error("unknown puzzle")]
+    UnknownPuzzle,
 }

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

@@ -1,9 +1,18 @@
 use chia_protocol::Bytes32;
 use chia_sdk_types::{MerkleTree, Mod, Vault1ofNArgs, VaultMofNArgs, VaultNofNArgs};
 use clvm_utils::TreeHash;
+use clvmr::NodePtr;
+
+use crate::{DriverError, SpendContext};
 
 use super::{KnownPuzzles, Member, PuzzleWithRestrictions, VaultLayer};
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum MofNOptimization {
+    One,
+    All,
+}
+
 #[derive(Debug, Clone)]
 pub struct MofN {
     required: usize,
@@ -26,6 +35,16 @@ impl MofN {
         &self.members
     }
 
+    fn optimization(&self) -> Option<MofNOptimization> {
+        if self.required == 1 {
+            Some(MofNOptimization::One)
+        } else if self.required == self.members.len() {
+            Some(MofNOptimization::All)
+        } else {
+            None
+        }
+    }
+
     fn merkle_tree(&self) -> MerkleTree {
         let leaves: Vec<Bytes32> = self
             .members
@@ -38,15 +57,40 @@ impl MofN {
 
 impl VaultLayer for MofN {
     fn puzzle_hash(&self) -> TreeHash {
-        if self.required == 1 {
-            let merkle_tree = self.merkle_tree();
-            Vault1ofNArgs::new(merkle_tree.root()).curry_tree_hash()
-        } else if self.required == self.members.len() {
-            let members = self.members.iter().map(VaultLayer::puzzle_hash).collect();
-            VaultNofNArgs::new(members).curry_tree_hash()
-        } else {
-            let merkle_tree = self.merkle_tree();
-            VaultMofNArgs::new(self.required, merkle_tree.root()).curry_tree_hash()
+        match self.optimization() {
+            Some(MofNOptimization::One) => {
+                let merkle_tree = self.merkle_tree();
+                Vault1ofNArgs::new(merkle_tree.root()).curry_tree_hash()
+            }
+            Some(MofNOptimization::All) => {
+                let members = self.members.iter().map(VaultLayer::puzzle_hash).collect();
+                VaultNofNArgs::new(members).curry_tree_hash()
+            }
+            None => {
+                let merkle_tree = self.merkle_tree();
+                VaultMofNArgs::new(self.required, merkle_tree.root()).curry_tree_hash()
+            }
+        }
+    }
+
+    fn puzzle(&self, ctx: &mut SpendContext) -> Result<NodePtr, DriverError> {
+        match self.optimization() {
+            Some(MofNOptimization::One) => {
+                let merkle_tree = self.merkle_tree();
+                ctx.curry(Vault1ofNArgs::new(merkle_tree.root()))
+            }
+            Some(MofNOptimization::All) => {
+                let members = self
+                    .members
+                    .iter()
+                    .map(|member| member.puzzle(ctx))
+                    .collect::<Result<_, _>>()?;
+                ctx.curry(VaultNofNArgs::new(members))
+            }
+            None => {
+                let merkle_tree = self.merkle_tree();
+                ctx.curry(VaultMofNArgs::new(self.required, merkle_tree.root()))
+            }
         }
     }
 

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

@@ -1,6 +1,9 @@
 use chia_bls::PublicKey;
 use chia_sdk_types::{BlsMember, Mod};
 use clvm_utils::TreeHash;
+use clvmr::NodePtr;
+
+use crate::{DriverError, SpendContext};
 
 use super::{KnownPuzzles, MofN, VaultLayer};
 
@@ -50,6 +53,14 @@ impl VaultLayer for Member {
         self.puzzle_hash
     }
 
+    fn puzzle(&self, ctx: &mut SpendContext) -> Result<NodePtr, DriverError> {
+        match &self.kind {
+            MemberKind::Bls(bls) => ctx.curry(bls),
+            MemberKind::MofN(m_of_n) => m_of_n.puzzle(ctx),
+            MemberKind::Unknown => Err(DriverError::UnknownPuzzle),
+        }
+    }
+
     fn replace(self, known_puzzles: &KnownPuzzles) -> Self {
         let kind = known_puzzles
             .members

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

@@ -1,5 +1,8 @@
 use chia_sdk_types::{DelegatedFeederArgs, IndexWrapperArgs, Mod, RestrictionsArgs};
 use clvm_utils::TreeHash;
+use clvmr::NodePtr;
+
+use crate::{DriverError, SpendContext};
 
 use super::{KnownPuzzles, Restriction, VaultLayer};
 
@@ -62,6 +65,35 @@ where
         IndexWrapperArgs::new(self.nonce, puzzle_hash).curry_tree_hash()
     }
 
+    fn puzzle(&self, ctx: &mut SpendContext) -> Result<NodePtr, DriverError> {
+        let mut puzzle = self.puzzle.puzzle(ctx)?;
+
+        if !self.restrictions.is_empty() {
+            let mut member_validators = Vec::new();
+            let mut delegated_puzzle_validators = Vec::new();
+
+            for restriction in &self.restrictions {
+                if restriction.is_member_condition_validator() {
+                    member_validators.push(restriction.puzzle(ctx)?);
+                } else {
+                    delegated_puzzle_validators.push(restriction.puzzle(ctx)?);
+                }
+            }
+
+            puzzle = ctx.curry(RestrictionsArgs::new(
+                member_validators,
+                delegated_puzzle_validators,
+                puzzle,
+            ))?;
+        }
+
+        if self.has_delegated_feeder {
+            puzzle = ctx.curry(DelegatedFeederArgs::new(puzzle))?;
+        }
+
+        ctx.curry(IndexWrapperArgs::new(self.nonce, puzzle))
+    }
+
     fn replace(mut self, known_puzzles: &KnownPuzzles) -> Self {
         let mut restrictions = Vec::with_capacity(self.restrictions.len());
         for restriction in self.restrictions {

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

@@ -1,5 +1,8 @@
 use chia_sdk_types::Timelock;
 use clvm_utils::TreeHash;
+use clvmr::NodePtr;
+
+use crate::{DriverError, SpendContext};
 
 use super::{KnownPuzzles, VaultLayer};
 
@@ -43,6 +46,13 @@ impl VaultLayer for Restriction {
         self.puzzle_hash
     }
 
+    fn puzzle(&self, ctx: &mut SpendContext) -> Result<NodePtr, DriverError> {
+        match &self.kind {
+            RestrictionKind::Timelock(timelock) => ctx.curry(timelock),
+            RestrictionKind::Unknown => Err(DriverError::UnknownPuzzle),
+        }
+    }
+
     fn replace(self, known_puzzles: &KnownPuzzles) -> Self {
         let kind = known_puzzles
             .restrictions

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

@@ -1,9 +1,13 @@
 use clvm_utils::TreeHash;
+use clvmr::NodePtr;
+
+use crate::{DriverError, SpendContext};
 
 use super::KnownPuzzles;
 
 pub trait VaultLayer {
     #[must_use]
     fn replace(self, known_puzzles: &KnownPuzzles) -> Self;
     fn puzzle_hash(&self) -> TreeHash;
+    fn puzzle(&self, ctx: &mut SpendContext) -> Result<NodePtr, DriverError>;
 }