ergotree_interpreter/sigma_protocol/
proof_tree.rs

1//! ProofTree
2
3extern crate derive_more;
4
5use std::fmt::Debug;
6
7use derive_more::From;
8use derive_more::TryInto;
9use ergotree_ir::sigma_protocol::sigma_boolean::SigmaBoolean;
10use ergotree_ir::sigma_protocol::sigma_boolean::SigmaConjectureItems;
11
12use crate::sigma_protocol::unproven_tree::CandUnproven;
13use crate::sigma_protocol::unproven_tree::UnprovenConjecture;
14use crate::sigma_protocol::UncheckedSchnorr;
15use crate::sigma_protocol::UnprovenSchnorr;
16
17use super::challenge::Challenge;
18use super::prover::ProverError;
19use super::unchecked_tree::UncheckedConjecture;
20use super::unchecked_tree::UncheckedDhTuple;
21use super::unchecked_tree::UncheckedTree;
22use super::unproven_tree::CorUnproven;
23use super::unproven_tree::CthresholdUnproven;
24use super::unproven_tree::UnprovenDhTuple;
25use super::unproven_tree::UnprovenLeaf;
26use super::unproven_tree::UnprovenTree;
27use super::FirstProverMessage;
28
29/// Proof tree
30#[derive(PartialEq, Debug, Clone, From, TryInto)]
31pub(crate) enum ProofTree {
32    /// Unchecked tree
33    UncheckedTree(UncheckedTree),
34    /// Unproven tree
35    UnprovenTree(UnprovenTree),
36}
37
38impl ProofTree {
39    /// Create a new proof tree with a new challenge
40    pub(crate) fn with_challenge(&self, challenge: Challenge) -> ProofTree {
41        match self {
42            ProofTree::UncheckedTree(uc) => uc.clone().with_challenge(challenge).into(),
43            ProofTree::UnprovenTree(ut) => ut.clone().with_challenge(challenge).into(),
44        }
45    }
46
47    pub(crate) fn as_tree_kind(&self) -> ProofTreeKind {
48        match self {
49            ProofTree::UncheckedTree(unch) => unch.as_tree_kind(),
50            ProofTree::UnprovenTree(unp) => unp.as_tree_kind(),
51        }
52    }
53
54    pub(crate) fn challenge(&self) -> Option<Challenge> {
55        match self {
56            ProofTree::UncheckedTree(unch) => Some(unch.challenge()),
57            ProofTree::UnprovenTree(unp) => unp.challenge(),
58        }
59    }
60}
61
62impl From<UncheckedSchnorr> for ProofTree {
63    fn from(v: UncheckedSchnorr) -> Self {
64        UncheckedTree::UncheckedLeaf(v.into()).into()
65    }
66}
67
68impl From<UncheckedDhTuple> for ProofTree {
69    fn from(v: UncheckedDhTuple) -> Self {
70        UncheckedTree::UncheckedLeaf(v.into()).into()
71    }
72}
73
74impl From<UnprovenSchnorr> for ProofTree {
75    fn from(v: UnprovenSchnorr) -> Self {
76        UnprovenTree::UnprovenLeaf(v.into()).into()
77    }
78}
79
80impl From<UnprovenDhTuple> for ProofTree {
81    fn from(v: UnprovenDhTuple) -> Self {
82        UnprovenTree::UnprovenLeaf(v.into()).into()
83    }
84}
85
86impl From<CandUnproven> for ProofTree {
87    fn from(v: CandUnproven) -> Self {
88        UnprovenTree::UnprovenConjecture(v.into()).into()
89    }
90}
91
92impl From<CorUnproven> for ProofTree {
93    fn from(v: CorUnproven) -> Self {
94        UnprovenTree::UnprovenConjecture(v.into()).into()
95    }
96}
97
98impl From<UnprovenConjecture> for ProofTree {
99    fn from(v: UnprovenConjecture) -> Self {
100        UnprovenTree::UnprovenConjecture(v).into()
101    }
102}
103
104impl From<UnprovenLeaf> for ProofTree {
105    fn from(v: UnprovenLeaf) -> Self {
106        UnprovenTree::UnprovenLeaf(v).into()
107    }
108}
109
110impl From<UncheckedConjecture> for ProofTree {
111    fn from(v: UncheckedConjecture) -> Self {
112        UncheckedTree::UncheckedConjecture(v).into()
113    }
114}
115
116impl From<CthresholdUnproven> for ProofTree {
117    fn from(v: CthresholdUnproven) -> Self {
118        UnprovenTree::UnprovenConjecture(v.into()).into()
119    }
120}
121
122/// Proof tree leaf
123pub trait ProofTreeLeaf: Debug {
124    /// Get proposition
125    fn proposition(&self) -> SigmaBoolean;
126
127    /// Get commitment
128    fn commitment_opt(&self) -> Option<FirstProverMessage>;
129}
130
131pub(crate) enum ConjectureType {
132    And = 0,
133    Or = 1,
134    Threshold = 2,
135}
136
137pub(crate) trait ProofTreeConjecture {
138    fn conjecture_type(&self) -> ConjectureType;
139    fn children(&self) -> SigmaConjectureItems<ProofTree>;
140}
141
142pub(crate) enum ProofTreeKind<'a> {
143    Leaf(&'a dyn ProofTreeLeaf),
144    Conjecture(&'a dyn ProofTreeConjecture),
145}
146
147/// Traverses the tree in the bottom-up manner, calling `f` for every node/leaf and setting
148/// it's returned value (if `Some`) as new node/leaf or do nothing if it's returned `None`
149pub(crate) fn rewrite_bu<F>(tree: ProofTree, f: &F) -> Result<ProofTree, ProverError>
150where
151    F: Fn(&ProofTree) -> Result<Option<ProofTree>, ProverError>,
152{
153    let cast_to_ust = |children: SigmaConjectureItems<ProofTree>| {
154        children.try_mapped(|c| {
155            if let ProofTree::UncheckedTree(ust) = c {
156                Ok(ust)
157            } else {
158                Err(ProverError::Unexpected(
159                    "rewrite: expected UncheckedSigmaTree got UnprovenTree",
160                ))
161            }
162        })
163    };
164
165    let tree_with_updated_children = match &tree {
166        ProofTree::UnprovenTree(unp_tree) => match unp_tree {
167            UnprovenTree::UnprovenLeaf(_) => tree,
168            UnprovenTree::UnprovenConjecture(conj) => match conj {
169                UnprovenConjecture::CandUnproven(cand) => UnprovenTree::UnprovenConjecture(
170                    UnprovenConjecture::CandUnproven(CandUnproven {
171                        children: cand.children.clone().try_mapped(|c| rewrite_bu(c, f))?,
172                        ..cand.clone()
173                    }),
174                )
175                .into(),
176                UnprovenConjecture::CorUnproven(cor) => {
177                    UnprovenTree::UnprovenConjecture(UnprovenConjecture::CorUnproven(CorUnproven {
178                        children: cor.children.clone().try_mapped(|c| rewrite_bu(c, f))?,
179                        ..cor.clone()
180                    }))
181                    .into()
182                }
183                UnprovenConjecture::CthresholdUnproven(ct) => {
184                    UnprovenTree::UnprovenConjecture(UnprovenConjecture::CthresholdUnproven(
185                        ct.clone()
186                            .with_children(ct.clone().children.try_mapped(|c| rewrite_bu(c, f))?),
187                    ))
188                    .into()
189                }
190            },
191        },
192        ProofTree::UncheckedTree(unch_tree) => match unch_tree {
193            UncheckedTree::UncheckedLeaf(_) => tree,
194            UncheckedTree::UncheckedConjecture(conj) => match conj {
195                UncheckedConjecture::CandUnchecked {
196                    challenge,
197                    children,
198                } => {
199                    let rewritten_children =
200                        children.clone().try_mapped(|c| rewrite_bu(c.into(), f))?;
201                    let casted_children = cast_to_ust(rewritten_children)?;
202                    UncheckedConjecture::CandUnchecked {
203                        children: casted_children,
204                        challenge: challenge.clone(),
205                    }
206                    .into()
207                }
208                UncheckedConjecture::CorUnchecked {
209                    challenge,
210                    children,
211                } => {
212                    let rewritten_children =
213                        children.clone().try_mapped(|c| rewrite_bu(c.into(), f))?;
214                    let casted_children = cast_to_ust(rewritten_children)?;
215                    UncheckedConjecture::CorUnchecked {
216                        children: casted_children,
217                        challenge: challenge.clone(),
218                    }
219                    .into()
220                }
221                UncheckedConjecture::CthresholdUnchecked {
222                    challenge,
223                    children,
224                    k,
225                    polynomial: polynomial_opt,
226                } => {
227                    let rewritten_children =
228                        children.clone().try_mapped(|c| rewrite_bu(c.into(), f))?;
229                    let casted_children = cast_to_ust(rewritten_children)?;
230                    UncheckedConjecture::CthresholdUnchecked {
231                        children: casted_children,
232                        challenge: challenge.clone(),
233                        k: *k,
234                        polynomial: polynomial_opt.clone(),
235                    }
236                    .into()
237                }
238            },
239        },
240    };
241    Ok(f(&tree_with_updated_children)?.unwrap_or(tree_with_updated_children))
242}
243
244/// Traverses the tree in the top-down manner, calling `f` for every node/leaf and setting
245/// it's returned value (if `Some`) as new node/leaf or do nothing if it's returned `None`
246pub(crate) fn rewrite_td<F>(tree: ProofTree, f: &F) -> Result<ProofTree, ProverError>
247where
248    F: Fn(&ProofTree) -> Result<Option<ProofTree>, ProverError>,
249{
250    let cast_to_ust = |children: SigmaConjectureItems<ProofTree>| {
251        children.try_mapped(|c| {
252            if let ProofTree::UncheckedTree(ust) = c {
253                Ok(ust)
254            } else {
255                Err(ProverError::Unexpected(
256                    "rewrite: expected UncheckedSigmaTree got UnprovenTree",
257                ))
258            }
259        })
260    };
261
262    let rewritten_tree = f(&tree)?.unwrap_or(tree);
263    Ok(match &rewritten_tree {
264        ProofTree::UnprovenTree(unp_tree) => match unp_tree {
265            UnprovenTree::UnprovenLeaf(_) => rewritten_tree,
266            UnprovenTree::UnprovenConjecture(conj) => match conj {
267                UnprovenConjecture::CandUnproven(cand) => UnprovenTree::UnprovenConjecture(
268                    UnprovenConjecture::CandUnproven(CandUnproven {
269                        children: cand.children.clone().try_mapped(|c| rewrite_td(c, f))?,
270                        ..cand.clone()
271                    }),
272                )
273                .into(),
274                UnprovenConjecture::CorUnproven(cor) => {
275                    UnprovenTree::UnprovenConjecture(UnprovenConjecture::CorUnproven(CorUnproven {
276                        children: cor.children.clone().try_mapped(|c| rewrite_td(c, f))?,
277                        ..cor.clone()
278                    }))
279                    .into()
280                }
281                UnprovenConjecture::CthresholdUnproven(ct) => {
282                    UnprovenTree::UnprovenConjecture(UnprovenConjecture::CthresholdUnproven(
283                        ct.clone()
284                            .with_children(ct.clone().children.try_mapped(|c| rewrite_td(c, f))?),
285                    ))
286                    .into()
287                }
288            },
289        },
290        ProofTree::UncheckedTree(unch_tree) => match unch_tree {
291            UncheckedTree::UncheckedLeaf(_) => rewritten_tree,
292            UncheckedTree::UncheckedConjecture(conj) => match conj {
293                UncheckedConjecture::CandUnchecked {
294                    challenge,
295                    children,
296                } => {
297                    let rewritten_children =
298                        children.clone().try_mapped(|c| rewrite_td(c.into(), f))?;
299                    let casted_children = cast_to_ust(rewritten_children)?;
300                    UncheckedConjecture::CandUnchecked {
301                        children: casted_children,
302                        challenge: challenge.clone(),
303                    }
304                    .into()
305                }
306                UncheckedConjecture::CorUnchecked {
307                    challenge,
308                    children,
309                } => {
310                    let rewritten_children =
311                        children.clone().try_mapped(|c| rewrite_td(c.into(), f))?;
312                    let casted_children = cast_to_ust(rewritten_children)?;
313                    UncheckedConjecture::CorUnchecked {
314                        children: casted_children,
315                        challenge: challenge.clone(),
316                    }
317                    .into()
318                }
319                UncheckedConjecture::CthresholdUnchecked {
320                    challenge,
321                    children,
322                    k,
323                    polynomial: polynomial_opt,
324                } => {
325                    let rewritten_children =
326                        children.clone().try_mapped(|c| rewrite_td(c.into(), f))?;
327                    let casted_children = cast_to_ust(rewritten_children)?;
328                    UncheckedConjecture::CthresholdUnchecked {
329                        children: casted_children,
330                        challenge: challenge.clone(),
331                        k: *k,
332                        polynomial: polynomial_opt.clone(),
333                    }
334                    .into()
335                }
336            },
337        },
338    })
339}