Skip to main content

cml_chain/builders/
vote_builder.rs

1use crate::{
2    crypto::hash::hash_plutus_data,
3    governance::{GovActionId, Voter, VotingProcedure, VotingProcedures},
4    plutus::PlutusData,
5    transaction::NativeScript,
6    RequiredSigners,
7};
8
9use super::{
10    utils::required_wits_from_required_signers,
11    witness_builder::{
12        InputAggregateWitnessData, NativeScriptWitnessInfo, PartialPlutusWitness,
13        RequiredWitnessSet,
14    },
15};
16
17#[derive(Debug, thiserror::Error)]
18pub enum VoteBuilderError {
19    #[error("Voter is script. Call with_plutus_vote() instead.")]
20    VoterIsScript,
21    #[error("Voter is key hash. Call with_vote() instead.")]
22    VoterIsKeyHash,
23    #[error("Vote already exists")]
24    VoteAlreayExists,
25    #[error("Missing the following witnesses for the input: {0:?}")]
26    MissingWitnesses(Box<RequiredWitnessSet>),
27}
28
29#[derive(Clone, Debug, Default)]
30pub struct VoteBuilderResult {
31    pub votes: VotingProcedures,
32    pub required_wits: RequiredWitnessSet,
33    pub aggregate_witnesses: Vec<InputAggregateWitnessData>,
34}
35
36#[derive(Clone, Debug)]
37pub struct VoteBuilder {
38    result: VoteBuilderResult,
39}
40
41impl Default for VoteBuilder {
42    fn default() -> Self {
43        Self::new()
44    }
45}
46
47impl VoteBuilder {
48    pub fn new() -> Self {
49        Self {
50            result: VoteBuilderResult::default(),
51        }
52    }
53
54    /// Add a vote using a voter with a key hash
55    /// Will throw an error if the voter is script-hash based
56    pub fn with_vote(
57        mut self,
58        voter: Voter,
59        gov_action_id: GovActionId,
60        procedure: VotingProcedure,
61    ) -> Result<Self, VoteBuilderError> {
62        if let Some(key_hash) = voter.key_hash() {
63            self.result.required_wits.add_vkey_key_hash(*key_hash);
64        } else {
65            return Err(VoteBuilderError::VoterIsScript);
66        }
67        if self
68            .result
69            .votes
70            .entry(voter)
71            .or_default()
72            .insert(gov_action_id, procedure)
73            .is_some()
74        {
75            return Err(VoteBuilderError::VoteAlreayExists);
76        }
77        Ok(self)
78    }
79
80    pub fn with_native_script_vote(
81        mut self,
82        voter: Voter,
83        gov_action_id: GovActionId,
84        procedure: VotingProcedure,
85        native_script: NativeScript,
86        witness_info: NativeScriptWitnessInfo,
87    ) -> Result<Self, VoteBuilderError> {
88        if let Some(script_hash) = voter.script_hash() {
89            if *script_hash != native_script.hash() {
90                let mut err_req_wits = RequiredWitnessSet::new();
91                err_req_wits.add_script_hash(*script_hash);
92                return Err(VoteBuilderError::MissingWitnesses(Box::new(err_req_wits)));
93            }
94            self.result.required_wits.add_script_hash(*script_hash);
95        } else {
96            return Err(VoteBuilderError::VoterIsKeyHash);
97        }
98
99        if self
100            .result
101            .votes
102            .entry(voter)
103            .or_default()
104            .insert(gov_action_id, procedure)
105            .is_some()
106        {
107            return Err(VoteBuilderError::VoteAlreayExists);
108        }
109
110        self.result
111            .aggregate_witnesses
112            .push(InputAggregateWitnessData::NativeScript(
113                native_script,
114                witness_info,
115            ));
116
117        Ok(self)
118    }
119
120    pub fn with_plutus_vote(
121        self,
122        voter: Voter,
123        gov_action_id: GovActionId,
124        procedure: VotingProcedure,
125        partial_witness: PartialPlutusWitness,
126        required_signers: RequiredSigners,
127        datum: PlutusData,
128    ) -> Result<Self, VoteBuilderError> {
129        self.with_plutus_vote_impl(
130            voter,
131            gov_action_id,
132            procedure,
133            partial_witness,
134            required_signers,
135            Some(datum),
136        )
137    }
138
139    pub fn with_plutus_vote_inline_datum(
140        self,
141        voter: Voter,
142        gov_action_id: GovActionId,
143        procedure: VotingProcedure,
144        partial_witness: PartialPlutusWitness,
145        required_signers: RequiredSigners,
146    ) -> Result<Self, VoteBuilderError> {
147        self.with_plutus_vote_impl(
148            voter,
149            gov_action_id,
150            procedure,
151            partial_witness,
152            required_signers,
153            None,
154        )
155    }
156
157    fn with_plutus_vote_impl(
158        mut self,
159        voter: Voter,
160        gov_action_id: GovActionId,
161        procedure: VotingProcedure,
162        partial_witness: PartialPlutusWitness,
163        required_signers: RequiredSigners,
164        datum: Option<PlutusData>,
165    ) -> Result<Self, VoteBuilderError> {
166        let mut required_wits = required_wits_from_required_signers(&required_signers);
167        if let Some(script_hash) = voter.script_hash() {
168            required_wits.add_script_hash(*script_hash);
169        } else {
170            return Err(VoteBuilderError::VoterIsKeyHash);
171        }
172
173        let mut required_wits_left = required_wits.clone();
174
175        // no way to know these at this time
176        required_wits_left.vkeys.clear();
177
178        let script_hash = partial_witness.script.hash();
179
180        // check the user provided all the required witnesses
181        required_wits_left.scripts.remove(&script_hash);
182        if let Some(datum) = &datum {
183            required_wits_left
184                .plutus_data
185                .remove(&hash_plutus_data(datum));
186        }
187
188        if required_wits_left.len() > 0 {
189            return Err(VoteBuilderError::MissingWitnesses(Box::new(
190                required_wits_left,
191            )));
192        }
193
194        if self
195            .result
196            .votes
197            .entry(voter)
198            .or_default()
199            .insert(gov_action_id, procedure)
200            .is_some()
201        {
202            return Err(VoteBuilderError::VoteAlreayExists);
203        }
204
205        self.result.required_wits.add_all(required_wits);
206
207        self.result
208            .aggregate_witnesses
209            .push(InputAggregateWitnessData::PlutusScript(
210                partial_witness,
211                required_signers,
212                datum,
213            ));
214
215        Ok(self)
216    }
217
218    pub fn build(self) -> VoteBuilderResult {
219        self.result
220    }
221}