From fbc61a7bbe75c99b29cd238f552c18542deb5c32 Mon Sep 17 00:00:00 2001
From: Haniel Barbosa <hanielbbarbosa@gmail.com>
Date: Sat, 3 Aug 2019 10:40:36 -0500
Subject: [PATCH] Collapse @ chains in SMT2 printer (#3140)

---
 src/options/printer_options.toml  |  8 ++++++++
 src/printer/smt2/smt2_printer.cpp | 30 +++++++++++++++++++++++++++++-
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/src/options/printer_options.toml b/src/options/printer_options.toml
index 89ae3b559..c1871259c 100644
--- a/src/options/printer_options.toml
+++ b/src/options/printer_options.toml
@@ -21,3 +21,11 @@ header = "options/printer_options.h"
   handler    = "stringToInstFormatMode"
   includes   = ["options/printer_modes.h"]
   help       = "print format mode for instantiations, see --inst-format=help"
+
+[[option]]
+  name       = "flattenHOChains"
+  category   = "regular"
+  long       = "flatten-ho-chains"
+  type       = "bool"
+  default    = "false"
+  help       = "print (binary) application chains in a flattened way, e.g. (a b c) rather than ((a b) c)"
diff --git a/src/printer/smt2/smt2_printer.cpp b/src/printer/smt2/smt2_printer.cpp
index 163c8acad..fd8b4ce85 100644
--- a/src/printer/smt2/smt2_printer.cpp
+++ b/src/printer/smt2/smt2_printer.cpp
@@ -24,6 +24,7 @@
 #include "expr/node_manager_attributes.h"
 #include "options/bv_options.h"
 #include "options/language.h"
+#include "options/printer_options.h"
 #include "options/smt_options.h"
 #include "printer/dagification_visitor.h"
 #include "smt/smt_engine.h"
@@ -509,7 +510,34 @@ void Smt2Printer::toStream(std::ostream& out,
   // uf theory
   case kind::APPLY_UF: typeChildren = true; break;
   // higher-order
-  case kind::HO_APPLY: break;
+  case kind::HO_APPLY:
+    if (!options::flattenHOChains())
+    {
+      break;
+    }
+    // collapse "@" chains, i.e.
+    //
+    // ((a b) c) --> (a b c)
+    //
+    // (((a b) ((c d) e)) f) --> (a b (c d e) f)
+    {
+      Node head = n;
+      std::vector<Node> args;
+      while (head.getKind() == kind::HO_APPLY)
+      {
+        args.insert(args.begin(), head[1]);
+        head = head[0];
+      }
+      toStream(out, head, toDepth, types, TypeNode::null());
+      for (unsigned i = 0, size = args.size(); i < size; ++i)
+      {
+        out << " ";
+        toStream(out, args[i], toDepth, types, TypeNode::null());
+      }
+      out << ")";
+    }
+    return;
+
   case kind::LAMBDA:
     out << smtKindString(k, d_variant) << " ";
     break;
-- 
2.30.2