Shell: Implement formatting for Heredocs

Userland/Shell/Formatter.cpp

@@ -7,6 +7,7 @@
 #include "Formatter.h"
 #include "AST.h"
 #include "Parser.h"
+#include <AK/Hex.h>
 #include <AK/ScopedValueRollback.h>
 #include <AK/TemporaryChange.h>
 
@@ -36,12 +37,14 @@ String Formatter::format()
 
     node->visit(*this);
 
-    auto string = m_builder.string_view();
+    VERIFY(m_builders.size() == 1);
+
+    auto string = current_builder().string_view();
 
     if (!string.ends_with(" "))
-        m_builder.append(m_trivia);
+        current_builder().append(m_trivia);
 
-    return m_builder.to_string();
+    return current_builder().to_string();
 }
 
 void Formatter::with_added_indent(int indent, Function<void()> callback)
@@ -63,6 +66,13 @@ void Formatter::in_new_block(Function<void()> callback)
     current_builder().append('}');
 }
 
+String Formatter::in_new_builder(Function<void()> callback, StringBuilder new_builder)
+{
+    m_builders.append(move(new_builder));
+    callback();
+    return m_builders.take_last().to_string();
+}
+
 void Formatter::test_and_update_output_cursor(const AST::Node* node)
 {
     if (!node)
@@ -96,9 +106,17 @@ void Formatter::will_visit(const AST::Node* node)
     }
 }
 
-void Formatter::insert_separator()
+void Formatter::insert_separator(bool escaped)
 {
+    if (escaped)
+        current_builder().append('\\');
     current_builder().append('\n');
+    if (!escaped && !m_heredocs_to_append_after_sequence.is_empty()) {
+        for (auto& entry : m_heredocs_to_append_after_sequence) {
+            current_builder().append(entry);
+        }
+        m_heredocs_to_append_after_sequence.clear();
+    }
     insert_indent();
 }
 
@@ -127,8 +145,8 @@ void Formatter::visit(const AST::And* node)
     with_added_indent(should_indent ? 1 : 0, [&] {
         node->left()->visit(*this);
 
-        current_builder().append(" \\");
+        current_builder().append(' ');
-        insert_separator();
+        insert_separator(true);
         current_builder().append("&& ");
 
         node->right()->visit(*this);
@@ -269,14 +287,17 @@ void Formatter::visit(const AST::DoubleQuotedString* node)
 {
     will_visit(node);
     test_and_update_output_cursor(node);
-    current_builder().append("\"");
+    auto not_in_heredoc = m_parent_node->kind() != AST::Node::Kind::Heredoc;
+    if (not_in_heredoc)
+        current_builder().append("\"");
 
     TemporaryChange quotes { m_options.in_double_quotes, true };
     TemporaryChange<const AST::Node*> parent { m_parent_node, node };
 
     NodeVisitor::visit(node);
 
-    current_builder().append("\"");
+    if (not_in_heredoc)
+        current_builder().append("\"");
     visited(node);
 }
 
@@ -355,6 +376,39 @@ void Formatter::visit(const AST::Glob* node)
     visited(node);
 }
 
+void Formatter::visit(const AST::Heredoc* node)
+{
+    will_visit(node);
+    test_and_update_output_cursor(node);
+
+    current_builder().append("<<");
+    if (node->deindent())
+        current_builder().append('~');
+    else
+        current_builder().append('-');
+
+    if (node->allow_interpolation())
+        current_builder().appendff("{}", node->end());
+    else
+        current_builder().appendff("'{}'", node->end());
+
+    auto content = in_new_builder([&] {
+        if (!node->contents())
+            return;
+
+        TemporaryChange<const AST::Node*> parent { m_parent_node, node };
+        TemporaryChange heredoc { m_options.in_heredoc, true };
+
+        auto& contents = *node->contents();
+        contents.visit(*this);
+        current_builder().appendff("\n{}\n", node->end());
+    });
+
+    m_heredocs_to_append_after_sequence.append(move(content));
+
+    visited(node);
+}
+
 void Formatter::visit(const AST::HistoryEvent* node)
 {
     will_visit(node);
@@ -567,8 +621,8 @@ void Formatter::visit(const AST::Or* node)
     with_added_indent(should_indent ? 1 : 0, [&] {
         node->left()->visit(*this);
 
-        current_builder().append(" \\");
+        current_builder().append(" ");
-        insert_separator();
+        insert_separator(true);
         current_builder().append("|| ");
 
         node->right()->visit(*this);
@@ -584,10 +638,10 @@ void Formatter::visit(const AST::Pipe* node)
     TemporaryChange<const AST::Node*> parent { m_parent_node, node };
 
     node->left()->visit(*this);
-    current_builder().append(" \\");
+    current_builder().append(" ");
 
     with_added_indent(should_indent ? 1 : 0, [&] {
-        insert_separator();
+        insert_separator(true);
         current_builder().append("| ");
 
         node->right()->visit(*this);
@@ -720,10 +774,10 @@ void Formatter::visit(const AST::StringLiteral* node)
 {
     will_visit(node);
     test_and_update_output_cursor(node);
-    if (!m_options.in_double_quotes)
+    if (!m_options.in_double_quotes && !m_options.in_heredoc)
         current_builder().append("'");
 
-    if (m_options.in_double_quotes) {
+    if (m_options.in_double_quotes && !m_options.in_heredoc) {
         for (auto ch : node->text()) {
             switch (ch) {
             case '"':
@@ -761,7 +815,7 @@ void Formatter::visit(const AST::StringLiteral* node)
         current_builder().append(node->text());
     }
 
-    if (!m_options.in_double_quotes)
+    if (!m_options.in_double_quotes && !m_options.in_heredoc)
         current_builder().append("'");
     visited(node);
 }

Userland/Shell/Formatter.h

@@ -8,6 +8,7 @@
 
 #include "NodeVisitor.h"
 #include <AK/Forward.h>
+#include <AK/String.h>
 #include <AK/StringBuilder.h>
 #include <AK/StringView.h>
 #include <AK/Types.h>
@@ -18,7 +19,7 @@ namespace Shell {
 class Formatter final : public AST::NodeVisitor {
 public:
     Formatter(const StringView& source, ssize_t cursor = -1)
-        : m_builder(round_up_to_power_of_two(source.length(), 16))
+        : m_builders({ StringBuilder { round_up_to_power_of_two(source.length(), 16) } })
         , m_source(source)
         , m_cursor(cursor)
     {
@@ -30,7 +31,8 @@ public:
     }
 
     explicit Formatter(const AST::Node& node)
-        : m_cursor(-1)
+        : m_builders({ StringBuilder {} })
+        , m_cursor(-1)
         , m_root_node(node)
     {
     }
@@ -57,6 +59,7 @@ private:
     virtual void visit(const AST::FunctionDeclaration*) override;
     virtual void visit(const AST::ForLoop*) override;
     virtual void visit(const AST::Glob*) override;
+    virtual void visit(const AST::Heredoc*) override;
     virtual void visit(const AST::HistoryEvent*) override;
     virtual void visit(const AST::Execute*) override;
     virtual void visit(const AST::IfCond*) override;
@@ -85,24 +88,26 @@ private:
     void test_and_update_output_cursor(const AST::Node*);
     void visited(const AST::Node*);
     void will_visit(const AST::Node*);
-    void insert_separator();
+    void insert_separator(bool escaped = false);
     void insert_indent();
 
     ALWAYS_INLINE void with_added_indent(int indent, Function<void()>);
     ALWAYS_INLINE void in_new_block(Function<void()>);
+    ALWAYS_INLINE String in_new_builder(Function<void()>, StringBuilder new_builder = StringBuilder {});
 
-    StringBuilder& current_builder() { return m_builder; }
+    StringBuilder& current_builder() { return m_builders.last(); }
 
     struct Options {
         size_t max_line_length_hint { 80 };
         bool explicit_parentheses { false };
         bool explicit_braces { false };
         bool in_double_quotes { false };
+        bool in_heredoc { false };
     } m_options;
 
     size_t m_current_indent { 0 };
 
-    StringBuilder m_builder;
+    Vector<StringBuilder> m_builders;
 
     StringView m_source;
     size_t m_output_cursor { 0 };
@@ -114,6 +119,7 @@ private:
     const AST::Node* m_last_visited_node { nullptr };
 
     StringView m_trivia;
+    Vector<String> m_heredocs_to_append_after_sequence;
 };
 
 }