factor out SExpr/SExprWriter classes out of smtlib backend, and also tidy them up/document them

Makefile

@@ -617,6 +617,7 @@ $(eval $(call add_include_file,kernel/register.h))
 $(eval $(call add_include_file,kernel/rtlil.h))
 $(eval $(call add_include_file,kernel/satgen.h))
 $(eval $(call add_include_file,kernel/scopeinfo.h))
+$(eval $(call add_include_file,kernel/sexpr.h))
 $(eval $(call add_include_file,kernel/sigtools.h))
 $(eval $(call add_include_file,kernel/timinginfo.h))
 $(eval $(call add_include_file,kernel/utils.h))
@@ -638,7 +639,7 @@ $(eval $(call add_include_file,backends/rtlil/rtlil_backend.h))
 
 OBJS += kernel/driver.o kernel/register.o kernel/rtlil.o kernel/log.o kernel/calc.o kernel/yosys.o
 OBJS += kernel/binding.o
-OBJS += kernel/cellaigs.o kernel/celledges.o kernel/cost.o kernel/satgen.o kernel/scopeinfo.o kernel/qcsat.o kernel/mem.o kernel/ffmerge.o kernel/ff.o kernel/yw.o kernel/json.o kernel/fmt.o
+OBJS += kernel/cellaigs.o kernel/celledges.o kernel/cost.o kernel/satgen.o kernel/scopeinfo.o kernel/qcsat.o kernel/mem.o kernel/ffmerge.o kernel/ff.o kernel/yw.o kernel/json.o kernel/fmt.o kernel/sexpr.o
 OBJS += kernel/drivertools.o kernel/functionalir.o
 ifeq ($(ENABLE_ZLIB),1)
 OBJS += kernel/fstdata.o

backends/functional/smtlib.cc

@@ -19,11 +19,14 @@
 
 #include "kernel/functionalir.h"
 #include "kernel/yosys.h"
+#include "kernel/sexpr.h"
 #include <ctype.h>
 
 USING_YOSYS_NAMESPACE
 PRIVATE_NAMESPACE_BEGIN
 
+using SExprUtil::list;
+
 const char *reserved_keywords[] = {
 	"BINARY", "DECIMAL", "HEXADECIMAL", "NUMERAL", "STRING", "_", "!", "as", "let", "exists", "forall", "match", "par",
 	"assert", "check-sat", "check-sat-assuming", "declare-const", "declare-datatype", "declare-datatypes",
@@ -47,166 +50,6 @@ struct SmtScope : public FunctionalTools::Scope<int> {
 	}
 };
 
-class SExpr {
-public:
-	std::variant<std::vector<SExpr>, std::string> _v;
-public:
-	SExpr(std::string a) : _v(std::move(a)) {}
-    SExpr(const char *a) : _v(a) {}
-	SExpr(int n) : _v(std::to_string(n)) {}
-	SExpr(std::vector<SExpr> const &l) : _v(l) {}
-	SExpr(std::vector<SExpr> &&l) : _v(std::move(l)) {}
-    bool is_atom() const { return std::holds_alternative<std::string>(_v); }
-    std::string const &atom() const { return std::get<std::string>(_v); }
-    bool is_list() const { return std::holds_alternative<std::vector<SExpr>>(_v); }
-    std::vector<SExpr> const &list() const { return std::get<std::vector<SExpr>>(_v); }
-	friend std::ostream &operator<<(std::ostream &os, SExpr const &sexpr) {
-		if(sexpr.is_atom())
-			os << sexpr.atom();
-		else if(sexpr.is_list()){
-			os << "(";
-			auto l = sexpr.list();
-			for(size_t i = 0; i < l.size(); i++) {
-				if(i > 0) os << " ";
-				os << l[i];
-			}
-			os << ")";
-		}else
-			os << "<invalid>";
-		return os;
-	}
-	std::string to_string() const {
-		std::stringstream ss;
-		ss << *this;
-		return ss.str();
-	}
-};
-template<typename... Args> SExpr list(Args&&... args) {
-	return SExpr(std::vector<SExpr>{std::forward<Args>(args)...});
-}
-
-class SExprWriter {
-    std::ostream &os;
-    int _max_line_width;
-    int _indent = 0;
-    int _pos = 0;
-    bool _pending_nl = false;
-	vector<bool> _unclosed;
-	vector<size_t> _unclosed_stack;
-	void nl_if_pending() {
-		if(_pending_nl) {
-            os << '\n';
-            _pos = 0;
-            _pending_nl = false;
-		}
-	}
-    void puts(std::string const &s) {
-        if(s.empty()) return;
-		nl_if_pending();
-        for(auto c : s) {
-            if(c == '\n') {
-                os << c;
-                _pos = 0;
-            } else {
-                if(_pos == 0) {
-                    for(int i = 0; i < _indent; i++)
-                        os << "  ";
-                    _pos = 2 * _indent;
-                }
-                os << c;
-                _pos++;
-            }
-        }
-    }
-    int width(SExpr const &sexpr) {
-        if(sexpr.is_atom())
-            return sexpr.atom().size();
-        else if(sexpr.is_list()) {
-            int w = 2;
-            for(auto arg : sexpr.list())
-                w += width(arg);
-            if(sexpr.list().size() > 1)
-                w += sexpr.list().size() - 1;
-            return w;
-        } else
-            return 0;
-    }
-    void print(SExpr const &sexpr, bool close = true, bool indent_rest = true) {
-        if(sexpr.is_atom())
-            puts(sexpr.atom());
-        else if(sexpr.is_list()) {
-            auto args = sexpr.list();
-			puts("(");
-			bool vertical = args.size() > 1 && _pos + width(sexpr) > _max_line_width;
-			if(vertical) _indent++;
-			for(size_t i = 0; i < args.size(); i++) {
-				if(i > 0) puts(vertical ? "\n" : " ");
-				print(args[i]);
-			}
-			_indent += (!close && indent_rest) - vertical;
-			if(close)
-				puts(")");
-			else {
-				_unclosed.push_back(indent_rest);
-				_pending_nl = true;
-			}
-        }else
-            log_error("shouldn't happen in SExprWriter::print");
-    }
-public:
-    SExprWriter(std::ostream &os, int max_line_width = 80)
-        : os(os)
-        , _max_line_width(max_line_width)
-    {}
-    void open(SExpr const &sexpr, bool indent_rest = true) {
-        log_assert(sexpr.is_list());
-        print(sexpr, false, indent_rest);
-    }
-    void close(size_t n = 1) {
-		log_assert(_unclosed.size() - (_unclosed_stack.empty() ? 0 : _unclosed_stack.back()) >= n);
-		while(n-- > 0) {
-			bool indented = _unclosed[_unclosed.size() - 1];
-			_unclosed.pop_back();
-			_pending_nl = _pos >= _max_line_width;
-			if(indented)
-				_indent--;
-			puts(")");
-			_pending_nl = true;
-		}
-    }
-	void push() {
-		_unclosed_stack.push_back(_unclosed.size());
-	}
-	void pop() {
-		auto t = _unclosed_stack.back();
-		log_assert(_unclosed.size() >= t);
-		close(_unclosed.size() - t);
-		_unclosed_stack.pop_back();
-	}
-    SExprWriter &operator <<(SExpr const &sexpr) {
-        print(sexpr);
-        _pending_nl = true;
-        return *this;
-    }
-	void comment(std::string const &str, bool hanging = false) {
-		if(hanging) {
-			if(_pending_nl) {
-				_pending_nl = false;
-				puts(" ");
-			}
-		}
-		puts("; ");
-		puts(str);
-		puts("\n");
-	}
-    ~SExprWriter() {
-		while(!_unclosed_stack.empty())
-			pop();
-		close(_unclosed.size());
-		nl_if_pending();
-    }
-};
-
 struct SmtSort {
 	FunctionalIR::Sort sort;
 	SmtSort(FunctionalIR::Sort sort) : sort(sort) {}

kernel/sexpr.cc

@@ -0,0 +1,163 @@
+/*
+ *  yosys -- Yosys Open SYnthesis Suite
+ *
+ *  Copyright (C) 2024  Emily Schmidt <emily@yosyshq.com>
+ *
+ *  Permission to use, copy, modify, and/or distribute this software for any
+ *  purpose with or without fee is hereby granted, provided that the above
+ *  copyright notice and this permission notice appear in all copies.
+ *
+ *  THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+ *  WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+ *  MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+ *  ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+ *  WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+ *  ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+ *  OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+ *
+ */
+
+#include "sexpr.h"
+
+YOSYS_NAMESPACE_BEGIN
+
+std::ostream &operator<<(std::ostream &os, SExpr const &sexpr) {
+    if(sexpr.is_atom())
+        os << sexpr.atom();
+    else if(sexpr.is_list()){
+        os << "(";
+        auto l = sexpr.list();
+        for(size_t i = 0; i < l.size(); i++) {
+            if(i > 0) os << " ";
+            os << l[i];
+        }
+        os << ")";
+    }else
+        os << "<invalid>";
+    return os;
+}
+
+ std::string SExpr::to_string() const {
+    std::stringstream ss;
+    ss << *this;
+    return ss.str();
+}
+
+void SExprWriter::nl_if_pending() {
+    if(_pending_nl) {
+        os << '\n';
+        _pos = 0;
+        _pending_nl = false;
+    }
+}
+
+void SExprWriter::puts(std::string_view s) {
+    if(s.empty()) return;
+    nl_if_pending();
+    for(auto c : s) {
+        if(c == '\n') {
+            os << c;
+            _pos = 0;
+        } else {
+            if(_pos == 0) {
+                for(int i = 0; i < _indent; i++)
+                    os << "  ";
+                _pos = 2 * _indent;
+            }
+            os << c;
+            _pos++;
+        }
+    }
+}
+
+
+// Calculate how much space would be left if expression was written
+// out in full horizontally. Returns any negative value if it doesn't fit.
+//
+// (Ideally we would avoid recalculating the widths of subexpression,
+// but I can't figure out how to store the widths. As an alternative,
+// we bail out of the calculation as soon as we can tell the expression
+// doesn't fit in the available space.)
+int SExprWriter::check_fit(SExpr const &sexpr, int space) {
+    if(sexpr.is_atom())
+        return space - sexpr.atom().size();
+    else if(sexpr.is_list()) {
+        space -= 2;
+        if(sexpr.list().size() > 1)
+            space -= sexpr.list().size() - 1;
+        for(auto arg : sexpr.list()) {
+            if(space < 0) break;
+            space = check_fit(arg, space);
+        }
+        return space;
+    } else
+        return -1;
+}
+
+void SExprWriter::print(SExpr const &sexpr, bool close, bool indent_rest) {
+    if(sexpr.is_atom())
+        puts(sexpr.atom());
+    else if(sexpr.is_list()) {
+        auto args = sexpr.list();
+        puts("(");
+        // Expressions are printed horizontally if they fit on the line.
+        // We do the check *after* puts("(") to make sure that _pos is accurate.
+        // (Otherwise there could be a pending newline + indentation)
+        bool vertical = args.size() > 1 && check_fit(sexpr, _max_line_width - _pos + 1) < 0;
+        if(vertical) _indent++;
+        for(size_t i = 0; i < args.size(); i++) {
+            if(i > 0) puts(vertical ? "\n" : " ");
+            print(args[i]);
+        }
+        // Any remaining arguments are currently always printed vertically,
+        // but are not indented if indent_rest = false.
+        _indent += (!close && indent_rest) - vertical;
+        if(close)
+            puts(")");
+        else {
+            _unclosed.push_back(indent_rest);
+            _pending_nl = true;
+        }
+    }else
+        log_error("shouldn't happen: SExpr '%s' is neither an atom nor a list", sexpr.to_string().c_str());
+}
+
+void SExprWriter::close(size_t n) {
+    log_assert(_unclosed.size() - (_unclosed_stack.empty() ? 0 : _unclosed_stack.back()) >= n);
+    while(n-- > 0) {
+        bool indented = _unclosed[_unclosed.size() - 1];
+        _unclosed.pop_back();
+        // Only print ) on the same line if it fits.
+        _pending_nl = _pos >= _max_line_width;
+        if(indented)
+            _indent--;
+        puts(")");
+        _pending_nl = true;
+    }
+}
+
+void SExprWriter::comment(std::string const &str, bool hanging) {
+    if(hanging) {
+        if(_pending_nl) {
+            _pending_nl = false;
+            puts(" ");
+        }
+    }
+    size_t i = 0, e;
+    do{
+        e = str.find('\n', i);
+        puts("; ");
+        puts(std::string_view(str).substr(i, e - i));
+        puts("\n");
+        i = e + 1;
+    }while(e != std::string::npos);
+}
+
+SExprWriter::~SExprWriter() {
+    while(!_unclosed_stack.empty())
+        pop();
+    close(_unclosed.size());
+    nl_if_pending();
+}
+
+YOSYS_NAMESPACE_END

kernel/sexpr.h

@@ -0,0 +1,122 @@
+/*
+ *  yosys -- Yosys Open SYnthesis Suite
+ *
+ *  Copyright (C) 2024  Emily Schmidt <emily@yosyshq.com>
+ *
+ *  Permission to use, copy, modify, and/or distribute this software for any
+ *  purpose with or without fee is hereby granted, provided that the above
+ *  copyright notice and this permission notice appear in all copies.
+ *
+ *  THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+ *  WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+ *  MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+ *  ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+ *  WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+ *  ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+ *  OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+ *
+ */
+
+#ifndef SEXPR_H
+#define SEXPR_H
+
+#include "kernel/yosys.h"
+
+YOSYS_NAMESPACE_BEGIN
+
+class SExpr {
+public:
+	std::variant<std::vector<SExpr>, std::string> _v;
+public:
+	SExpr(std::string a) : _v(std::move(a)) {}
+    SExpr(const char *a) : _v(a) {}
+    // FIXME: should maybe be defined for all integral types
+	SExpr(int n) : _v(std::to_string(n)) {}
+	SExpr(std::vector<SExpr> const &l) : _v(l) {}
+	SExpr(std::vector<SExpr> &&l) : _v(std::move(l)) {}
+    // It would be nicer to have an std::initializer_list constructor,
+    // but that causes confusing issues with overload resolution sometimes.
+    template<typename... Args> static SExpr list(Args&&... args) {
+	    return SExpr(std::vector<SExpr>{std::forward<Args>(args)...});
+    }
+    bool is_atom() const { return std::holds_alternative<std::string>(_v); }
+    std::string const &atom() const { return std::get<std::string>(_v); }
+    bool is_list() const { return std::holds_alternative<std::vector<SExpr>>(_v); }
+    std::vector<SExpr> const &list() const { return std::get<std::vector<SExpr>>(_v); }
+	std::string to_string() const;
+};
+
+std::ostream &operator<<(std::ostream &os, SExpr const &sexpr);
+
+namespace SExprUtil {
+    // A little hack so that `using SExprUtil::list` lets you import a shortcut to `SExpr::list`
+    template<typename... Args> SExpr list(Args&&... args) {
+	    return SExpr(std::vector<SExpr>{std::forward<Args>(args)...});
+    }
+}
+
+// SExprWriter is a pretty printer for s-expr. It does not try very hard to get a good layout.
+class SExprWriter {
+    std::ostream &os;
+    int _max_line_width;
+    int _indent = 0;
+    int _pos = 0;
+    // If _pending_nl is set, print a newline before the next character.
+    // This lets us "undo" the last newline so we can put
+    // closing parentheses or a hanging comment on the same line.
+    bool _pending_nl = false;
+    // Unclosed parentheses (boolean stored is indent_rest)
+	vector<bool> _unclosed;
+    // Used only for push() and pop() (stores _unclosed.size())
+	vector<size_t> _unclosed_stack;
+	void nl_if_pending();
+    void puts(std::string_view s);
+    int check_fit(SExpr const &sexpr, int space);
+    void print(SExpr const &sexpr, bool close = true, bool indent_rest = true);
+public:
+    SExprWriter(std::ostream &os, int max_line_width = 80)
+        : os(os)
+        , _max_line_width(max_line_width)
+    {}
+    // Print an s-expr.
+    SExprWriter &operator <<(SExpr const &sexpr) {
+        print(sexpr);
+        _pending_nl = true;
+        return *this;
+    }
+    // Print an s-expr (which must be a list), but leave room for extra elements
+    // which may be printed using either << or further calls to open.
+    // If indent_rest = false, the remaining elements are not intended
+    // (for avoiding unreasonable indentation on deeply nested structures).
+    void open(SExpr const &sexpr, bool indent_rest = true) {
+        log_assert(sexpr.is_list());
+        print(sexpr, false, indent_rest);
+    }
+    // Close the s-expr opened with the last call to open
+    // (if an argument is given, close that many s-exprs).
+    void close(size_t n = 1);
+    // push() remembers how many s-exprs are currently open
+	void push() {
+		_unclosed_stack.push_back(_unclosed.size());
+	}
+    // pop() closes all s-expr opened since the corresponding call to push()
+	void pop() {
+		auto t = _unclosed_stack.back();
+		log_assert(_unclosed.size() >= t);
+		close(_unclosed.size() - t);
+		_unclosed_stack.pop_back();
+	}
+    // Print a comment.
+    // If hanging = true, append it to the end of the last printed s-expr.
+	void comment(std::string const &str, bool hanging = false);
+    // Flush any unprinted characters to the std::ostream, but does not close unclosed parentheses.
+    void flush() {
+        nl_if_pending();
+    }
+    // Destructor closes any unclosed parentheses and flushes.
+    ~SExprWriter();
+};
+
+YOSYS_NAMESPACE_END
+
+#endif