Improve post-rewrite expression formatting (#472)

Co-authored-by: Gabe Jackson <gj@mail.co.de>

polar-core/src/formatting.rs

@@ -410,17 +410,23 @@ pub mod to_polar {
                         )
                     }
                 }
-                // `Dot` sometimes formats as a predicate
+                // Lookup operator
                 Dot => {
-                    if self.args.len() == 2 {
-                        let call_term = if let Value::String(s) = self.args[1].value() {
-                            s.to_string()
-                        } else {
-                            self.args[1].to_polar()
-                        };
-                        format!("{}.{}", self.args[0].to_polar(), call_term)
+                    let call_term = if let Value::String(s) = self.args[1].value() {
+                        s.to_string()
                     } else {
-                        format!(".({})", format_args(self.operator, &self.args, ", "))
+                        self.args[1].to_polar()
+                    };
+                    match self.args.len() {
+                        2 => format!("{}.{}", self.args[0].to_polar(), call_term),
+                        3 => format!(
+                            "{}.{} = {}",
+                            self.args[0].to_polar(),
+                            call_term,
+                            self.args[2].to_polar()
+                        ),
+                        // Invalid
+                        _ => format!(".({})", format_args(self.operator, &self.args, ", ")),
                     }
                 }
                 // Unary operators
@@ -431,12 +437,27 @@ pub mod to_polar {
                 ),
                 // Binary operators
                 Mul | Div | Add | Sub | Eq | Geq | Leq | Neq | Gt | Lt | Unify | Isa | In
-                | Assign => format!(
-                    "{} {} {}",
-                    to_polar_parens(self.operator, &self.args[0]),
-                    self.operator.to_polar(),
-                    to_polar_parens(self.operator, &self.args[1])
-                ),
+                | Assign => match self.args.len() {
+                    2 => format!(
+                        "{} {} {}",
+                        to_polar_parens(self.operator, &self.args[0]),
+                        self.operator.to_polar(),
+                        to_polar_parens(self.operator, &self.args[1]),
+                    ),
+                    3 => format!(
+                        "{} {} {} = {}",
+                        to_polar_parens(self.operator, &self.args[0]),
+                        self.operator.to_polar(),
+                        to_polar_parens(self.operator, &self.args[1]),
+                        to_polar_parens(self.operator, &self.args[2]),
+                    ),
+                    // Invalid
+                    _ => format!(
+                        "{}({})",
+                        self.operator.to_polar(),
+                        format_args(self.operator, &self.args, ", ")
+                    ),
+                },
                 // n-ary operators
                 And => format_args(
                     self.operator,

polar-core/src/rewrites.rs

@@ -215,26 +215,21 @@ mod tests {
 
         // First rewrite
         rewrite_rule(&mut rule, &mut kb);
-        assert_eq!(rule.to_polar(), "f(_value_1) if .(a, \"b\", _value_1);");
+        assert_eq!(rule.to_polar(), "f(_value_1) if a.b = _value_1;");
 
         // Check we can parse the rules back again
         let again = parse_rules(&rule.to_polar());
         let again_rule = again[0].clone();
         assert_eq!(again_rule.to_polar(), rule.to_polar());
 
-        // Call rewrite again
-        let mut rewrite_again_rule = again_rule.clone();
-        rewrite_rule(&mut rewrite_again_rule, &mut kb);
-        assert_eq!(rewrite_again_rule.to_polar(), again_rule.to_polar());
-
         // Chained lookups
         let rules = parse_rules("f(a.b.c);");
         let mut rule = rules[0].clone();
         assert_eq!(rule.to_polar(), "f(a.b.c);");
         rewrite_rule(&mut rule, &mut kb);
         assert_eq!(
             rule.to_polar(),
-            "f(_value_2) if .(a, \"b\", _value_3) and .(_value_3, \"c\", _value_2);"
+            "f(_value_2) if a.b = _value_3 and _value_3.c = _value_2;"
         );
     }
 
@@ -249,7 +244,7 @@ mod tests {
         rewrite_rule(&mut rule, &mut kb);
         assert_eq!(
             rule.to_polar(),
-            "f(a, c) if .(a, b(c), _value_1) and _value_1;"
+            "f(a, c) if a.b(c) = _value_1 and _value_1;"
         );
 
         // Nested lookups
@@ -259,7 +254,7 @@ mod tests {
         rewrite_rule(&mut rule, &mut kb);
         assert_eq!(
             rule.to_polar(),
-            "f(a, c, e) if .(e, f(), _value_4) and .(c, d(_value_4), _value_3) and .(a, b(_value_3), _value_2) and _value_2;"
+            "f(a, c, e) if e.f() = _value_4 and c.d(_value_4) = _value_3 and a.b(_value_3) = _value_2 and _value_2;"
         );
     }
 
@@ -269,26 +264,39 @@ mod tests {
         let mut term = parse_query("x and a.b");
         assert_eq!(term.to_polar(), "x and a.b");
         rewrite_term(&mut term, &mut kb);
-        assert_eq!(term.to_polar(), "x and .(a, \"b\", _value_1) and _value_1");
+        assert_eq!(term.to_polar(), "x and a.b = _value_1 and _value_1");
 
         let mut query = parse_query("f(a.b().c)");
         assert_eq!(query.to_polar(), "f(a.b().c)");
         rewrite_term(&mut query, &mut kb);
         assert_eq!(
             query.to_polar(),
-            ".(a, b(), _value_3) and .(_value_3, \"c\", _value_2) and f(_value_2)"
+            "a.b() = _value_3 and _value_3.c = _value_2 and f(_value_2)"
         );
 
         let mut term = parse_query("a.b = 1");
         rewrite_term(&mut term, &mut kb);
-        assert_eq!(term.to_polar(), ".(a, \"b\", _value_4) and _value_4 = 1");
+        assert_eq!(term.to_polar(), "a.b = _value_4 and _value_4 = 1");
         let mut term = parse_query("{x: 1}.x = 1");
         assert_eq!(term.to_polar(), "{x: 1}.x = 1");
         rewrite_term(&mut term, &mut kb);
-        assert_eq!(
-            term.to_polar(),
-            ".({x: 1}, \"x\", _value_5) and _value_5 = 1"
-        );
+        assert_eq!(term.to_polar(), "{x: 1}.x = _value_5 and _value_5 = 1");
+    }
+
+    #[test]
+    fn rewrite_expressions() {
+        let mut kb = KnowledgeBase::new();
+
+        let mut term = parse_query("0 - 0 = 0");
+        assert_eq!(term.to_polar(), "0 - 0 = 0");
+        rewrite_term(&mut term, &mut kb);
+        assert_eq!(term.to_polar(), "0 - 0 = _op_1 and _op_1 = 0");
+
+        let rules = parse_rules("sum(a, b, a + b);");
+        let mut rule = rules[0].clone();
+        assert_eq!(rule.to_polar(), "sum(a, b, a + b);");
+        rewrite_rule(&mut rule, &mut kb);
+        assert_eq!(rule.to_polar(), "sum(a, b, _op_2) if a + b = _op_2;");
     }
 
     #[test]
@@ -299,15 +307,15 @@ mod tests {
         rewrite_term(&mut term, &mut kb);
         assert_eq!(
             term.to_polar(),
-            ".(bar, \"y\", _value_2) and new (Foo(x: _value_2), _instance_1) and _instance_1"
+            "bar.y = _value_2 and new (Foo(x: _value_2), _instance_1) and _instance_1"
         );
 
         let mut term = parse_query("f(new Foo(x: bar.y))");
         assert_eq!(term.to_polar(), "f(new Foo(x: bar.y))");
         rewrite_term(&mut term, &mut kb);
         assert_eq!(
             term.to_polar(),
-            ".(bar, \"y\", _value_4) and new (Foo(x: _value_4), _instance_3) and f(_instance_3)"
+            "bar.y = _value_4 and new (Foo(x: _value_4), _instance_3) and f(_instance_3)"
         );
     }
 
@@ -358,9 +366,6 @@ mod tests {
         assert_eq!(term.to_polar(), "not foo.x = 1");
 
         rewrite_term(&mut term, &mut kb);
-        pretty_assertions::assert_eq!(
-            term.to_polar(),
-            "not (.(foo, \"x\", _value_1) and _value_1 = 1)"
-        )
+        pretty_assertions::assert_eq!(term.to_polar(), "not (foo.x = _value_1 and _value_1 = 1)")
     }
 }