Add SigType check in CompilerExtData for deciding key byte-lengths

SarcasticNastik · SarcasticNastik · commit b3825d3953a8 · 2022-04-07T02:46:18.000+05:30
diff --git a/src/policy/compiler.rs b/src/policy/compiler.rs
@@ -22,6 +22,7 @@ use std::convert::From;
 use std::marker::PhantomData;
 use std::{cmp, error, f64, fmt, mem};
 
+use miniscript::context::SigType;
 use miniscript::limits::MAX_PUBKEYS_PER_MULTISIG;
 use miniscript::types::{self, ErrorKind, ExtData, Property, Type};
 use miniscript::ScriptContext;
@@ -136,7 +137,7 @@ impl CompilationKey {
 }
 
 #[derive(Copy, Clone, Debug)]
-struct CompilerExtData {
+struct CompilerExtData<Ctx: ScriptContext> {
     /// If this node is the direct child of a disjunction, this field must
     /// have the probability of its branch being taken. Otherwise it is ignored.
     /// All functions initialize it to `None`.
@@ -148,14 +149,18 @@ struct CompilerExtData {
     /// (total length of all witness pushes, plus their own length prefixes)
     /// for fragments that can be dissatisfied without failing the script.
     dissat_cost: Option<f64>,
+    /// Signature Type for deciding the number of bytes required for [`CompilerExtData::sat_cost`]
+    /// and [`CompilerExtData::dissat_cost`] for the given script.
+    sig_type: PhantomData<Ctx>,
 }
 
-impl Property for CompilerExtData {
+impl<Ctx: ScriptContext> Property for CompilerExtData<Ctx> {
     fn from_true() -> Self {
         CompilerExtData {
             branch_prob: None,
             sat_cost: 0.0,
             dissat_cost: None,
+            sig_type: PhantomData,
         }
     }
 
@@ -164,38 +169,64 @@ impl Property for CompilerExtData {
             branch_prob: None,
             sat_cost: f64::MAX,
             dissat_cost: Some(0.0),
+            sig_type: PhantomData,
         }
     }
 
     fn from_pk_k() -> Self {
         CompilerExtData {
             branch_prob: None,
-            sat_cost: 73.0,
+            sat_cost: match Ctx::sig_type() {
+                SigType::Ecdsa => 73.,
+                SigType::Schnorr => 64.,
+            },
             dissat_cost: Some(1.0),
+            sig_type: PhantomData,
         }
     }
 
     fn from_pk_h() -> Self {
         CompilerExtData {
             branch_prob: None,
-            sat_cost: 73.0 + 34.0,
-            dissat_cost: Some(1.0 + 34.0),
+            sat_cost: match Ctx::sig_type() {
+                SigType::Ecdsa => 73.0 + 34.0,
+                SigType::Schnorr => 64. + 32.,
+            },
+            dissat_cost: Some(
+                1.0 + match Ctx::sig_type() {
+                    SigType::Ecdsa => 34.0,
+                    SigType::Schnorr => 32.,
+                },
+            ),
+            sig_type: PhantomData,
         }
     }
 
     fn from_multi(k: usize, _n: usize) -> Self {
         CompilerExtData {
             branch_prob: None,
-            sat_cost: 1.0 + 73.0 * k as f64,
+            sat_cost: 1.0
+                + match Ctx::sig_type() {
+                    SigType::Ecdsa => 73.0 * k as f64,
+                    SigType::Schnorr => 64. * k as f64,
+                },
             dissat_cost: Some(1.0 * (k + 1) as f64),
+            sig_type: PhantomData,
         }
     }
 
     fn from_hash() -> Self {
         CompilerExtData {
             branch_prob: None,
-            sat_cost: 33.0,
-            dissat_cost: Some(33.0),
+            sat_cost: match Ctx::sig_type() {
+                SigType::Ecdsa => 33.0,
+                SigType::Schnorr => 32.,
+            },
+            dissat_cost: Some(match Ctx::sig_type() {
+                SigType::Ecdsa => 33.0,
+                SigType::Schnorr => 32.,
+            }),
+            sig_type: PhantomData,
         }
     }
 
@@ -204,6 +235,7 @@ impl Property for CompilerExtData {
             branch_prob: None,
             sat_cost: 0.0,
             dissat_cost: None,
+            sig_type: PhantomData,
         }
     }
 
@@ -212,6 +244,7 @@ impl Property for CompilerExtData {
             branch_prob: None,
             sat_cost: self.sat_cost,
             dissat_cost: self.dissat_cost,
+            sig_type: PhantomData,
         })
     }
 
@@ -220,6 +253,7 @@ impl Property for CompilerExtData {
             branch_prob: None,
             sat_cost: self.sat_cost,
             dissat_cost: self.dissat_cost,
+            sig_type: PhantomData,
         })
     }
 
@@ -228,6 +262,7 @@ impl Property for CompilerExtData {
             branch_prob: None,
             sat_cost: self.sat_cost,
             dissat_cost: self.dissat_cost,
+            sig_type: PhantomData,
         })
     }
 
@@ -236,6 +271,7 @@ impl Property for CompilerExtData {
             branch_prob: None,
             sat_cost: 2.0 + self.sat_cost,
             dissat_cost: Some(1.0),
+            sig_type: PhantomData,
         })
     }
 
@@ -244,6 +280,7 @@ impl Property for CompilerExtData {
             branch_prob: None,
             sat_cost: self.sat_cost,
             dissat_cost: None,
+            sig_type: PhantomData,
         })
     }
 
@@ -252,6 +289,7 @@ impl Property for CompilerExtData {
             branch_prob: None,
             sat_cost: self.sat_cost,
             dissat_cost: Some(1.0),
+            sig_type: PhantomData,
         })
     }
 
@@ -260,6 +298,7 @@ impl Property for CompilerExtData {
             branch_prob: None,
             sat_cost: self.sat_cost,
             dissat_cost: self.dissat_cost,
+            sig_type: PhantomData,
         })
     }
 
@@ -268,6 +307,7 @@ impl Property for CompilerExtData {
             branch_prob: None,
             sat_cost: self.sat_cost,
             dissat_cost: None,
+            sig_type: PhantomData,
         })
     }
 
@@ -281,6 +321,7 @@ impl Property for CompilerExtData {
             branch_prob: None,
             sat_cost: 2.0 + self.sat_cost,
             dissat_cost: Some(1.0),
+            sig_type: PhantomData,
         })
     }
 
@@ -289,6 +330,7 @@ impl Property for CompilerExtData {
             branch_prob: None,
             sat_cost: 1.0 + self.sat_cost,
             dissat_cost: Some(2.0),
+            sig_type: PhantomData,
         })
     }
 
@@ -300,6 +342,7 @@ impl Property for CompilerExtData {
                 (Some(l), Some(r)) => Some(l + r),
                 _ => None,
             },
+            sig_type: PhantomData,
         })
     }
 
@@ -308,6 +351,16 @@ impl Property for CompilerExtData {
             branch_prob: None,
             sat_cost: left.sat_cost + right.sat_cost,
             dissat_cost: None,
+            sig_type: PhantomData,
+        })
+    }
+
+    fn and_n(a: Self, b: Self) -> Result<Self, types::ErrorKind> {
+        Ok(CompilerExtData {
+            branch_prob: None,
+            sat_cost: a.sat_cost + b.sat_cost,
+            dissat_cost: a.dissat_cost,
+            sig_type: PhantomData,
         })
     }
 
@@ -323,6 +376,7 @@ impl Property for CompilerExtData {
             sat_cost: lprob * (l.sat_cost + r.dissat_cost.unwrap())
                 + rprob * (r.sat_cost + l.dissat_cost.unwrap()),
             dissat_cost: Some(l.dissat_cost.unwrap() + r.dissat_cost.unwrap()),
+            sig_type: PhantomData,
         })
     }
 
@@ -337,6 +391,7 @@ impl Property for CompilerExtData {
             branch_prob: None,
             sat_cost: lprob * l.sat_cost + rprob * (r.sat_cost + l.dissat_cost.unwrap()),
             dissat_cost: r.dissat_cost.map(|rd| l.dissat_cost.unwrap() + rd),
+            sig_type: PhantomData,
         })
     }
 
@@ -351,6 +406,7 @@ impl Property for CompilerExtData {
             branch_prob: None,
             sat_cost: lprob * l.sat_cost + rprob * (r.sat_cost + l.dissat_cost.unwrap()),
             dissat_cost: None,
+            sig_type: PhantomData,
         })
     }
 
@@ -377,6 +433,7 @@ impl Property for CompilerExtData {
             } else {
                 None
             },
+            sig_type: PhantomData,
         })
     }
 
@@ -400,14 +457,7 @@ impl Property for CompilerExtData {
             } else {
                 None
             },
-        })
-    }
-
-    fn and_n(a: Self, b: Self) -> Result<Self, types::ErrorKind> {
-        Ok(CompilerExtData {
-            branch_prob: None,
-            sat_cost: a.sat_cost + b.sat_cost,
-            dissat_cost: a.dissat_cost,
+            sig_type: PhantomData,
         })
     }
 
@@ -427,6 +477,7 @@ impl Property for CompilerExtData {
             branch_prob: None,
             sat_cost: sat_cost * k_over_n + dissat_cost * (1.0 - k_over_n),
             dissat_cost: Some(dissat_cost),
+            sig_type: PhantomData,
         })
     }
 }
@@ -437,7 +488,7 @@ struct AstElemExt<Pk: MiniscriptKey, Ctx: ScriptContext> {
     /// The actual Miniscript fragment with type information
     ms: Arc<Miniscript<Pk, Ctx>>,
     /// Its "type" in terms of compiler data
-    comp_ext_data: CompilerExtData,
+    comp_ext_data: CompilerExtData<Ctx>,
 }
 
 impl<Pk: MiniscriptKey, Ctx: ScriptContext> AstElemExt<Pk, Ctx> {
@@ -470,8 +521,8 @@ impl<Pk: MiniscriptKey, Ctx: ScriptContext> AstElemExt<Pk, Ctx> {
         r: &AstElemExt<Pk, Ctx>,
     ) -> Result<AstElemExt<Pk, Ctx>, types::Error<Pk, Ctx>> {
         let lookup_ext = |n| match n {
-            0 => Some(l.comp_ext_data),
-            1 => Some(r.comp_ext_data),
+            0 => Some(l.comp_ext_data.clone()),
+            1 => Some(r.comp_ext_data.clone()),
             _ => unreachable!(),
         };
         //Types and ExtData are already cached and stored in children. So, we can
@@ -497,9 +548,9 @@ impl<Pk: MiniscriptKey, Ctx: ScriptContext> AstElemExt<Pk, Ctx> {
         c: &AstElemExt<Pk, Ctx>,
     ) -> Result<AstElemExt<Pk, Ctx>, types::Error<Pk, Ctx>> {
         let lookup_ext = |n| match n {
-            0 => Some(a.comp_ext_data),
-            1 => Some(b.comp_ext_data),
-            2 => Some(c.comp_ext_data),
+            0 => Some(a.comp_ext_data.clone()),
+            1 => Some(b.comp_ext_data.clone()),
+            2 => Some(c.comp_ext_data.clone()),
             _ => unreachable!(),
         };
         //Types and ExtData are already cached and stored in children. So, we can
@@ -525,7 +576,7 @@ struct Cast<Pk: MiniscriptKey, Ctx: ScriptContext> {
     node: fn(Arc<Miniscript<Pk, Ctx>>) -> Terminal<Pk, Ctx>,
     ast_type: fn(types::Type) -> Result<types::Type, ErrorKind>,
     ext_data: fn(types::ExtData) -> Result<types::ExtData, ErrorKind>,
-    comp_ext_data: fn(CompilerExtData) -> Result<CompilerExtData, types::ErrorKind>,
+    comp_ext_data: fn(CompilerExtData<Ctx>) -> Result<CompilerExtData<Ctx>, types::ErrorKind>,
 }
 
 impl<Pk: MiniscriptKey, Ctx: ScriptContext> Cast<Pk, Ctx> {
@@ -537,7 +588,7 @@ impl<Pk: MiniscriptKey, Ctx: ScriptContext> Cast<Pk, Ctx> {
                 node: (self.node)(Arc::clone(&ast.ms)),
                 phantom: PhantomData,
             }),
-            comp_ext_data: (self.comp_ext_data)(ast.comp_ext_data)?,
+            comp_ext_data: (self.comp_ext_data)(ast.comp_ext_data.clone())?,
         })
     }
 }
@@ -955,19 +1006,19 @@ where
                 let bw = best(types::Base::W, policy_cache, ast, sp, dp)?;
 
                 let diff = be.cost_1d(sp, dp) - bw.cost_1d(sp, dp);
-                best_es.push((be.comp_ext_data, be));
-                best_ws.push((bw.comp_ext_data, bw));
+                best_es.push((be.comp_ext_data.clone(), be));
+                best_ws.push((bw.comp_ext_data.clone(), bw));
 
                 if diff < min_value.1 {
                     min_value.0 = i;
                     min_value.1 = diff;
                 }
             }
-            sub_ext_data.push(best_es[min_value.0].0);
+            sub_ext_data.push(best_es[min_value.0].0.clone());
             sub_ast.push(Arc::clone(&best_es[min_value.0].1.ms));
             for (i, _ast) in subs.iter().enumerate() {
                 if i != min_value.0 {
-                    sub_ext_data.push(best_ws[i].0);
+                    sub_ext_data.push(best_ws[i].0.clone());
                     sub_ast.push(Arc::clone(&best_ws[i].1.ms));
                 }
             }
@@ -978,7 +1029,7 @@ where
                     Miniscript::from_ast(ast)
                         .expect("threshold subs, which we just compiled, typeck"),
                 ),
-                comp_ext_data: CompilerExtData::threshold(k, n, |i| Ok(sub_ext_data[i]))
+                comp_ext_data: CompilerExtData::threshold(k, n, |i| Ok(sub_ext_data[i].clone()))
                     .expect("threshold subs, which we just compiled, typeck"),
             };
             insert_wrap!(ast_ext);
