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};
 
@@ -60,7 +61,7 @@ pub enum Policy<Pk: MiniscriptKey> {
 /// 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)]
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
 pub struct PolicyDiff<Pk: MiniscriptKey> {
     /// The first policy
     pub a: Vec<Policy<Pk>>,
@@ -85,21 +86,49 @@ impl<Pk: MiniscriptKey> PolicyDiff<Pk> {
 
     /// Create a new PolicyDiff in the
     pub fn new(a: Policy<Pk>, b: Policy<Pk>) -> Self {
-        match (a.normalized(), b.normalized()) {
-            (x, y) if x == y => Self::_new(vec![], vec![]),
-            (Policy::Threshold(k1, subs1), Policy::Threshold(k2, subs2))
-                if k1 == k2 && subs1.len() == subs2.len() =>
-            {
-                let mut a = Self::_new(vec![], vec![]);
-
-                for (sub_a, sub_b) in subs1.into_iter().zip(subs2.into_iter()) {
-                    let sub_diff = Self::_new(vec![sub_a], vec![sub_b]);
-                    a.combine(sub_diff);
+        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)],
+                        )
+                    }
                 }
-                a
+                (x, y) => PolicyDiff::_new(vec![x], vec![y]),
             }
-            (x, y) => Self::_new(vec![x], vec![y]),
         }
+
+        new_helper(a.normalized(), b.normalized())
     }
 }
 
@@ -187,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
@@ -715,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