Implement fmt::Display for variable types

Author: uncomputable

src/core/types.rs

@@ -1,5 +1,6 @@
 use crate::merkle::common::{MerkleRoot, TypeMerkleRoot};
 use crate::merkle::tmr::Tmr;
+use crate::util::replace_to_buffer;
 use std::collections::VecDeque;
 use std::iter::FromIterator;
 use std::{cell::RefCell, cmp, fmt, rc::Rc, sync::Arc};
@@ -128,6 +129,65 @@ pub(crate) enum VariableType {
     Product(RcVar, RcVar),
 }
 
+impl VariableType {
+    fn uncompressed_display<W: fmt::Write>(&self, w: &mut W) -> fmt::Result {
+        match self {
+            VariableType::Unit => w.write_str("1"),
+            VariableType::Sum(a, b) => {
+                w.write_str("(")?;
+                a.borrow().inner.uncompressed_display(w)?;
+                w.write_str(" + ")?;
+                b.borrow().inner.uncompressed_display(w)?;
+                w.write_str(")")
+            }
+            VariableType::Product(a, b) => {
+                w.write_str("(")?;
+                a.borrow().inner.uncompressed_display(w)?;
+                w.write_str(" × ")?;
+                b.borrow().inner.uncompressed_display(w)?;
+                w.write_str(")")
+            }
+            // Uncomment to use intermediate compression: more allocation of smaller strings
+            // VariableType::Sum(a, b) => write!(w, "({} + {})", a, b),
+            // VariableType::Product(a, b) => write!(w, "({} × {})", a, b),
+        }
+    }
+}
+
+fn into_compressed_display(mut buffer1: String) -> String {
+    let mut buffer2 = String::with_capacity(buffer1.len());
+
+    replace_to_buffer(&buffer1, &mut buffer2, "(1 + 1)", "2");
+    buffer1.clear();
+    replace_to_buffer(&buffer2, &mut buffer1, "(2 × 2)", "2^2");
+    buffer2.clear();
+    replace_to_buffer(&buffer1, &mut buffer2, "(2^2 × 2^2)", "2^4");
+    buffer1.clear();
+    replace_to_buffer(&buffer2, &mut buffer1, "(2^4 × 2^4)", "2^8");
+    buffer2.clear();
+    replace_to_buffer(&buffer1, &mut buffer2, "(2^8 × 2^8)", "2^16");
+    buffer1.clear();
+    replace_to_buffer(&buffer2, &mut buffer1, "(2^16 × 2^16)", "2^32");
+    buffer2.clear();
+    replace_to_buffer(&buffer1, &mut buffer2, "(2^32 × 2^32)", "2^64");
+    buffer1.clear();
+    replace_to_buffer(&buffer2, &mut buffer1, "(2^64 × 2^64)", "2^128");
+    buffer2.clear();
+    replace_to_buffer(&buffer1, &mut buffer2, "(2^128 × 2^128)", "2^256");
+
+    buffer2
+}
+
+impl fmt::Display for VariableType {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        let mut buffer1 = String::new();
+        self.uncompressed_display(&mut buffer1)?;
+        let buffer2 = into_compressed_display(buffer1);
+
+        write!(f, "{}", buffer2)
+    }
+}
+
 /// Variable that can be cheaply cloned and internally mutated
 pub(crate) type RcVar = Rc<RefCell<Variable>>;
 
@@ -182,6 +242,16 @@ impl fmt::Debug for Variable {
     }
 }
 
+impl fmt::Display for Variable {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        let mut buffer1 = String::new();
+        self.inner.uncompressed_display(&mut buffer1)?;
+        let buffer2 = into_compressed_display(buffer1);
+
+        write!(f, "{}", buffer2)
+    }
+}
+
 /// Unification variable without metadata (see [`Variable`])
 #[derive(Clone)]
 pub(crate) enum VariableInner {
@@ -208,6 +278,19 @@ impl fmt::Debug for VariableInner {
     }
 }
 
+impl VariableInner {
+    fn uncompressed_display<W: fmt::Write>(&self, w: &mut W) -> fmt::Result {
+        match self {
+            VariableInner::Free(id) => w.write_str(id),
+            VariableInner::Bound(ty, _) => ty.uncompressed_display(w),
+            // Uncomment to use intermediate compression: more allocation of smaller strings
+            // VariableInner::Bound(ty, _) => write!(w, "{}", ty),
+            VariableInner::EqualTo(parent) => parent.borrow().inner.uncompressed_display(w),
+            VariableInner::Finalized(ty) => write!(w, "{}", ty),
+        }
+    }
+}
+
 /// Factory for creating free variables with fresh names.
 /// Identifiers are assigned sequentially as follows:
 /// `A`, `B`, `C`, ... `Z`, `AA`, `AB`, `AC`, ...

src/util.rs

@@ -65,3 +65,25 @@ pub(crate) fn bitvec_to_bytevec(bitvec: &[bool]) -> Vec<u8> {
     }
     ret
 }
+
+/// Replace all occurrences of a `pattern` in a `src` string with a `replacement` string,
+/// and write the result to a `dst` string.
+///
+/// This method does not allocate if `dst` is sufficiently long (e.g., as long as `src`).
+pub(crate) fn replace_to_buffer(src: &str, dst: &mut String, pattern: &str, replacement: &str) {
+    let mut last_index = 0;
+    let len = pattern.len();
+
+    for (index, _) in src.match_indices(pattern) {
+        if last_index < index {
+            dst.push_str(&src[last_index..index])
+        }
+
+        dst.push_str(replacement);
+        last_index = index + len;
+    }
+
+    if last_index < src.len() {
+        dst.push_str(&src[last_index..])
+    }
+}