Pretty print a policy

src/policy/semantic.rs

@@ -14,6 +14,7 @@
 
 //! Abstract Policies
 
+use std::collections::HashSet;
 use std::str::FromStr;
 use std::{fmt, str};
 
@@ -57,6 +58,80 @@ pub enum Policy<Pk: MiniscriptKey> {
     Threshold(usize, Vec<Policy<Pk>>),
 }
 
+/// Compute the Policy difference between two policies.
+/// This is useful when trying to find out the conditions
+/// under which the two policies are different.
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
+pub struct PolicyDiff<Pk: MiniscriptKey> {
+    /// The first policy
+    pub a: Vec<Policy<Pk>>,
+    /// The second policy
+    pub b: Vec<Policy<Pk>>,
+}
+
+impl<Pk: MiniscriptKey> PolicyDiff<Pk> {
+    /// Combine two policy differences.
+    // Policies should not generally contain repeated conditions,
+    // therefore we do not many attempt to deal with those.
+    pub fn combine(&mut self, second: Self) {
+        self.a.extend(second.a);
+        self.b.extend(second.b);
+    }
+
+    // create a new diff directly without removing
+    // same policies.
+    fn _new(a: Vec<Policy<Pk>>, b: Vec<Policy<Pk>>) -> Self {
+        Self { a, b }
+    }
+
+    /// Create a new PolicyDiff in the
+    pub fn new(a: Policy<Pk>, b: Policy<Pk>) -> Self {
+        pub fn new_helper<Pk: MiniscriptKey>(a: Policy<Pk>, b: Policy<Pk>) -> PolicyDiff<Pk> {
+            match (a, b) {
+                (ref x, ref y) if x == y => PolicyDiff::_new(vec![], vec![]),
+                (Policy::Threshold(k1, subs1), Policy::Threshold(k2, subs2)) => {
+                    if k1 == k2 && subs1.len() == subs2.len() {
+                        let mut ind_a = HashSet::new();
+                        let mut ind_b = HashSet::new();
+                        for i in 0..subs1.len() {
+                            let sub_a = &subs1[i];
+                            let b_pos = subs2.iter().position(|sub_b| sub_a == sub_b);
+                            match b_pos {
+                                Some(j) => {
+                                    ind_a.insert(i);
+                                    ind_b.insert(j);
+                                }
+                                None => {}
+                            }
+                        }
+                        let diff_a: Vec<_> = subs1
+                            .into_iter()
+                            .enumerate()
+                            .filter(|(i, _x)| !ind_a.contains(i))
+                            .map(|(_i, p)| p)
+                            .collect();
+                        let diff_b = subs2
+                            .into_iter()
+                            .enumerate()
+                            .filter(|(i, _x)| !ind_b.contains(i))
+                            .map(|(_i, p)| p)
+                            .collect::<Vec<_>>();
+                        PolicyDiff::_new(diff_a, diff_b)
+                    } else {
+                        PolicyDiff::_new(
+                            vec![Policy::Threshold(k1, subs1)],
+                            vec![Policy::Threshold(k2, subs2)],
+                        )
+                    }
+                }
+                (x, y) => PolicyDiff::_new(vec![x], vec![y]),
+            }
+        }
+
+        new_helper(a.normalized(), b.normalized())
+    }
+}
+
 impl<Pk: MiniscriptKey> ForEachKey<Pk> for Policy<Pk> {
     fn for_each_key<'a, F: FnMut(ForEach<'a, Pk>) -> bool>(&'a self, mut pred: F) -> bool
     where
@@ -78,6 +153,42 @@ impl<Pk: MiniscriptKey> ForEachKey<Pk> for Policy<Pk> {
 }
 
 impl<Pk: MiniscriptKey> Policy<Pk> {
+    fn term_name(&self) -> String {
+        match *self {
+            Policy::Threshold(k, ref subs) => {
+                if k == subs.len() {
+                    String::from(format!("and({},{})", k, subs.len()))
+                } else if k == 1 {
+                    String::from(format!("or({},{})", k, subs.len()))
+                } else {
+                    String::from(format!("thresh({},{})", k, subs.len()))
+                }
+            }
+            _ => self.to_string(),
+        }
+    }
+
+    fn _pprint_tree(&self, prefix: String, last: bool) {
+        let prefix_current = if last { "`-- " } else { "|-- " };
+
+        println!("{}{}{}", prefix, prefix_current, self.term_name());
+
+        let prefix_child = if last { "    " } else { "|   " };
+        let prefix = prefix + prefix_child;
+
+        if let Policy::Threshold(_k, ref subs) = *self {
+            let last_child = subs.len() - 1;
+
+            for (i, child) in subs.iter().enumerate() {
+                child._pprint_tree(prefix.to_string(), i == last_child);
+            }
+        }
+    }
+
+    /// Pretty Print a tree
+    pub fn pprint_tree(&self) {
+        self._pprint_tree("".to_string(), true);
+    }
     /// Convert a policy using one kind of public key to another
     /// type of public key
     pub fn translate_pkh<Fpkh, Q, E>(&self, mut translatefpkh: Fpkh) -> Result<Policy<Q>, E>
@@ -105,6 +216,67 @@ impl<Pk: MiniscriptKey> Policy<Pk> {
         }
     }
 
+    fn _pprint_diff(&self, other: &Policy<Pk>, prefix: String, last: bool) {
+        match (self, other) {
+            (x, y) if x == y => x._pprint_tree(prefix, last),
+            (Policy::Threshold(k1, subs1), Policy::Threshold(k2, subs2))
+                if k1 == k2 && subs1.len() == subs2.len() =>
+            {
+                let prefix_current = if last { "`-- " } else { "|-- " };
+
+                println!("{}{}{}", prefix, prefix_current, self.term_name());
+
+                let prefix_child = if last { "    " } else { "|   " };
+                let prefix = prefix + prefix_child;
+
+                let last_child = subs1.len() - 1;
+
+                // Hashsets to maintain which children of subs are matched
+                // onto children from others
+                let mut ind_a = HashSet::new();
+                let mut ind_b = HashSet::new();
+                for i in 0..subs1.len() {
+                    let sub_a = &subs1[i];
+                    let b_pos = subs2.iter().position(|sub_b| sub_a == sub_b);
+                    match b_pos {
+                        Some(j) => {
+                            ind_a.insert(i);
+                            ind_b.insert(j);
+                        }
+                        None => {}
+                    }
+                }
+                let mut j = 0;
+                for (i, sub_a) in subs1.iter().enumerate() {
+                    // The position in subs2
+                    if ind_a.contains(&i) {
+                        sub_a._pprint_tree(prefix.to_string(), last);
+                    } else {
+                        while ind_b.contains(&j) {
+                            j = j + 1;
+                        }
+                        sub_a._pprint_diff(&subs2[j], prefix.to_string(), i == last_child);
+                    }
+                }
+            }
+            (x, y) => {
+                let red = "\x1b[0;31m";
+                let nc = "\x1b[0m";
+                let green = "\x1b[0;32m";
+                print!("{}", green);
+                x._pprint_tree(prefix.clone(), last);
+                print!("{}{}", nc, red);
+                y._pprint_tree(prefix, last);
+                print!("{}", nc);
+            }
+        }
+    }
+
+    /// Pretty Print a tree
+    pub fn pprint_diff(&self, other: &Policy<Pk>) {
+        self._pprint_diff(other, "".to_string(), true);
+    }
+
     /// This function computes whether the current policy entails the second one.
     /// A |- B means every satisfaction of A is also a satisfaction of B.
     /// This implementation will run slow for larger policies but should be sufficient for
@@ -633,6 +805,34 @@ mod tests {
         );
     }
 
+    #[test]
+    fn policy_diff() {
+        let pol1 = StringPolicy::from_str("or(pkh(A),pkh(C))").unwrap();
+        let pol2 = StringPolicy::from_str("or(pkh(B),pkh(C))").unwrap();
+        let diff = PolicyDiff::new(pol1.clone(), pol2.clone());
+        assert_eq!(
+            diff,
+            PolicyDiff::new(
+                StringPolicy::from_str("pkh(A)").unwrap(),
+                StringPolicy::from_str("pkh(B)").unwrap()
+            )
+        );
+        // Uncomment for pprint
+        // pol1.pprint_diff(&pol2);
+
+        let pol1 = StringPolicy::from_str("or(pkh(A),pkh(C))").unwrap();
+        // change the order
+        let pol2 = StringPolicy::from_str("or(pkh(C),and(pkh(B),older(9)))").unwrap();
+        let diff = PolicyDiff::new(pol1.clone(), pol2.clone());
+        assert_eq!(
+            diff,
+            PolicyDiff::new(
+                StringPolicy::from_str("pkh(A)").unwrap(),
+                StringPolicy::from_str("and(pkh(B),older(9))").unwrap()
+            )
+        );
+    }
+
     #[test]
     fn entailment_liquid_test() {
         //liquid policy
@@ -641,6 +841,8 @@ mod tests {
         // Very bad idea to add master key,pk but let's have it have 50M blocks
         let master_key = StringPolicy::from_str("and(older(50000000),pkh(master))").unwrap();
         let new_liquid_pol = Policy::Threshold(1, vec![liquid_pol.clone(), master_key]);
+        // Pretty print a policy
+        // liquid_pol.pprint_tree();
 
         assert!(liquid_pol.clone().entails(new_liquid_pol.clone()).unwrap());
         assert!(!new_liquid_pol.entails(liquid_pol.clone()).unwrap());