Add simple newline insertion to formatter

This fixes "bracebug" (reported here: google#299) and many
variants of it.

Of course it's not a complete solution to any newline problems.
In particular it never removes or rearranges newlines, it only
adds some. But it still should help a lot.

core/ast.h

@@ -92,6 +92,8 @@ static inline std::string ASTTypeToString(ASTType type) {
         case AST_UNARY: return "AST_UNARY";
         case AST_VAR: return "AST_VAR";
     }
+    std::cerr << "Invalid AST type" << "\n";
+    abort();
 }
 
 /** Represents a variable / parameter / field name. */

core/formatter.cpp

@@ -963,6 +963,280 @@ class PrettyFieldNames : public FmtPass {
     }
 };
 
+/// Add newlines inside complex structures (arrays, objects etc.).
+///
+/// The main principle is that a structure can either be:
+/// * expanded and contain newlines in all the designated places
+/// * unexpanded and not contain any newlines in any of designated places.
+///
+/// It only looks shallowly at the AST nodes, so there may be some newlines deeper that
+/// don't affect expanding. For example:
+/// [{
+///     'a': 'b',
+///     'c': 'd',
+/// }]
+/// The outer array can be unexpanded, because there are no newlines between
+class FixNewlines: public FmtPass {
+    using FmtPass::visit;
+
+    bool shouldExpand(const Array *array)
+    {
+        for (const auto &elem: array->elements) {
+            if (countNewlines(open_fodder(elem.expr)) > 0) {
+                return true;
+            }
+        }
+        if (countNewlines(array->closeFodder) > 0) {
+            return true;
+        }
+        return false;
+    }
+
+    void expand(Array *array)
+    {
+        for (auto &elem: array->elements) {
+            ensureCleanNewline(open_fodder(elem.expr));
+        }
+        ensureCleanNewline(array->closeFodder);
+    }
+
+    Fodder &objectFieldOpenFodder(ObjectField &field)
+    {
+        if (field.kind == ObjectField::Kind::FIELD_STR) {
+            return field.expr1->openFodder;
+        }
+        return field.fodder1;
+    }
+
+    bool shouldExpand(Object *object)
+    {
+        for (auto &field: object->fields) {
+            if (countNewlines(objectFieldOpenFodder(field)) > 0) {
+                return true;
+            }
+        }
+        if (countNewlines(object->closeFodder) > 0) {
+            return true;
+        }
+        return false;
+    }
+
+    void expand(Object *object)
+    {
+        for (auto &field: object->fields) {
+            ensureCleanNewline(objectFieldOpenFodder(field));
+        }
+        ensureCleanNewline(object->closeFodder);
+    }
+
+    bool shouldExpand(Local *local)
+    {
+        for (auto &bind: local->binds) {
+            if (countNewlines(bind.varFodder) > 0) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    void expand(Local *local)
+    {
+        bool first = true;
+        for (auto &bind: local->binds) {
+            if (!first) {
+                ensureCleanNewline(bind.varFodder);
+            }
+            first = false;
+        }
+    }
+
+    bool shouldExpand(ArrayComprehension *comp)
+    {
+        if (countNewlines(open_fodder(comp->body)) > 0) {
+            return true;
+        }
+        for (auto &spec: comp->specs) {
+            if (countNewlines(spec.openFodder) > 0) {
+                return true;
+            }
+        }
+        if (countNewlines(comp->closeFodder) > 0) {
+            return true;
+        }
+        return false;
+    }
+
+    void expand(ArrayComprehension *comp)
+    {
+        ensureCleanNewline(open_fodder(comp->body));
+        for (auto &spec: comp->specs) {
+            ensureCleanNewline(spec.openFodder);
+        }
+        ensureCleanNewline(comp->closeFodder);
+    }
+
+    bool shouldExpand(ObjectComprehension *comp)
+    {
+        for (auto &field: comp->fields) {
+            if (countNewlines(objectFieldOpenFodder(field)) > 0) {
+                return true;
+            }
+        }
+        for (auto &spec: comp->specs) {
+            if (countNewlines(spec.openFodder) > 0) {
+                return true;
+            }
+        }
+        if (countNewlines(comp->closeFodder) > 0) {
+            return true;
+        }
+        return false;
+    }
+
+    void expand(ObjectComprehension *comp)
+    {
+        for (auto &field: comp->fields) {
+            ensureCleanNewline(objectFieldOpenFodder(field));
+        }
+        for (auto &spec: comp->specs) {
+            ensureCleanNewline(spec.openFodder);
+        }
+        ensureCleanNewline(comp->closeFodder);
+    }
+
+    bool shouldExpand(Parens *parens)
+    {
+        return countNewlines(open_fodder(parens->expr)) > 0
+            || countNewlines(parens->closeFodder) > 0;
+    }
+
+    void expand(Parens *parens)
+    {
+        ensureCleanNewline(open_fodder(parens->expr));
+        ensureCleanNewline(parens->closeFodder);
+    }
+
+    Fodder &argParamOpenFodder(ArgParam &param)
+    {
+        if (param.id != nullptr) {
+            return param.idFodder;
+        } else if (param.expr != nullptr) {
+            return open_fodder(param.expr);
+        } else {
+            std::cerr << "Invalid ArgParam" << std::endl;
+            abort();
+        }
+    }
+
+    // Example:
+    // f(1, 2,
+    //   3)
+    // Should be expanded to:
+    // f(1,
+    //   2,
+    //   3)
+    bool shouldExpandBetween(ArgParams &params)
+    {
+        bool first = true;
+        for (auto &param: params) {
+            if (!first && countNewlines(argParamOpenFodder(param)) > 0) {
+                return true;
+            }
+            first = false;
+        }
+        return false;
+    }
+
+    void expandBetween(ArgParams &params)
+    {
+        bool first = true;
+        for (auto &param: params) {
+            if (!first) {
+                ensureCleanNewline(argParamOpenFodder(param));
+            }
+            first = false;
+        }
+    }
+
+    // Example:
+    // very_long_function_on_a_long_line_so_arguments_dont_fit_on_the_same_line(
+    //     1, 2, 3)
+    // Should be expanded to:
+    // very_long_function_on_a_long_line_so_arguments_dont_fit_on_the_same_line(
+    //     1, 2, 3
+    // )
+    bool shouldExpandNearParens(ArgParams &params, Fodder &fodder_r)
+    {
+        if (params.empty()) {
+            return false;
+        }
+        auto &argFodder = argParamOpenFodder(params.front());
+        return countNewlines(fodder_r) > 0 || countNewlines(argFodder) > 0;
+    }
+
+    void expandNearParens(ArgParams &params, Fodder &fodder_r)
+    {
+        if (!params.empty()) {
+             ensureCleanNewline(argParamOpenFodder(params.front()));
+        }
+        ensureCleanNewline(fodder_r);
+    }
+
+    public:
+    FixNewlines(Allocator &alloc, const FmtOpts &opts) : FmtPass(alloc, opts) { }
+
+    template<class T>
+    void simpleExpandingVisit(T *expr) {
+        if (shouldExpand(expr)) {
+            expand(expr);
+        }
+        FmtPass::visit(expr);
+    }
+
+    void visit(Array *array)
+    {
+        simpleExpandingVisit(array);
+    }
+
+    void visit(Object *object)
+    {
+        simpleExpandingVisit(object);
+    }
+
+    void visit(Local *local)
+    {
+        simpleExpandingVisit(local);
+    }
+
+    void visit(ArrayComprehension *comp)
+    {
+        simpleExpandingVisit(comp);
+    }
+
+    void visit(ObjectComprehension *comp)
+    {
+        simpleExpandingVisit(comp);
+    }
+
+    void visit(Parens *parens)
+    {
+        simpleExpandingVisit(parens);
+    }
+
+    void params(Fodder &fodder_l, ArgParams &params, Fodder &fodder_r)
+    {
+        if (shouldExpandBetween(params)) {
+            expandBetween(params);
+        }
+
+        if (shouldExpandNearParens(params, fodder_r)) {
+            expandNearParens(params, fodder_r);
+        }
+
+        FmtPass::params(fodder_l, params, fodder_r);
+    }
+};
+
 class FixIndentation {
 
     FmtOpts opts;
@@ -977,7 +1251,7 @@ class FixIndentation {
      * \param separate_token If the last fodder was an interstitial, whether a space should follow
      * it.
      * \param all_but_last_indent New indentation value for all but final fodder element.
-     * \param last_indent New indentation value for but final fodder element.
+     * \param last_indent New indentation value for the final fodder element.
      */
     void fill(Fodder &fodder, bool space_before, bool separate_token,
               unsigned all_but_last_indent, unsigned last_indent)
@@ -1889,13 +2163,6 @@ class SortImports {
         return false;
     }
 
-    void ensureCleanNewline(Fodder &fodder)
-    {
-        if (!fodder_has_clean_endline(fodder)) {
-            fodder_push_back(fodder, FodderElement(FodderElement::Kind::LINE_END, 0, 0, {}));
-        }
-    }
-
     AST *toplevelImport(Local *local, ImportElems &imports, const Fodder &groupOpenFodder)
     {
         assert(isGoodLocal(local));
@@ -1958,6 +2225,7 @@ std::string jsonnet_fmt(AST *ast, Fodder &final_fodder, const FmtOpts &opts)
     remove_initial_newlines(ast);
     if (opts.maxBlankLines > 0)
         EnforceMaximumBlankLines(alloc, opts).file(ast, final_fodder);
+    FixNewlines(alloc, opts).file(ast, final_fodder);
     FixTrailingCommas(alloc, opts).file(ast, final_fodder);
     FixParens(alloc, opts).file(ast, final_fodder);
     FixPlusObject(alloc, opts).file(ast, final_fodder);

core/lexer.h

@@ -180,12 +180,42 @@ static inline void fodder_move_front(Fodder &a, Fodder &b)
     b.clear();
 }
 
-static inline Fodder make_fodder(const FodderElement &elem) {
+static inline Fodder make_fodder(const FodderElement &elem)
+{
     Fodder fodder;
     fodder_push_back(fodder, elem);
     return fodder;
 }
 
+static inline void ensureCleanNewline(Fodder &fodder)
+{
+    if (!fodder_has_clean_endline(fodder)) {
+        fodder_push_back(fodder, FodderElement(FodderElement::Kind::LINE_END, 0, 0, {}));
+    }
+}
+
+static inline int countNewlines(const FodderElement &elem)
+{
+    switch (elem.kind) {
+        case FodderElement::INTERSTITIAL:
+            return 0;
+        case FodderElement::LINE_END:
+            return 1;
+        case FodderElement::PARAGRAPH:
+            return elem.comment.size() + elem.blanks;
+    }
+    std::cerr << "Unknown FodderElement kind" << std::endl;
+    abort();
+}
+
+static inline int countNewlines(const Fodder &fodder)
+{
+    int sum = 0;
+    for (const auto &elem: fodder) {
+        sum += countNewlines(elem);
+    }
+    return sum;
+}
 
 static inline std::ostream &operator<<(std::ostream &o, const Fodder &fodder)
 {

core/pass.cpp

@@ -116,12 +116,7 @@ void CompilerPass::expr(AST *&ast_)
 void CompilerPass::visit(Apply *ast)
 {
     expr(ast->target);
-    fodder(ast->fodderL);
-    for (auto &arg : ast->args) {
-        expr(arg.expr);
-        fodder(arg.commaFodder);
-    }
-    fodder(ast->fodderR);
+    params(ast->fodderL, ast->args, ast->fodderR);
     if (ast->tailstrict) {
         fodder(ast->tailstrictFodder);
     }
@@ -232,7 +227,7 @@ void CompilerPass::visit(Local *ast)
         fodder(bind.varFodder);
         if (bind.functionSugar) {
             params(bind.parenLeftFodder, bind.params, bind.parenRightFodder);
-        } 
+        }
         fodder(bind.opFodder);
         expr(bind.body);
         fodder(bind.closeFodder);

examples/bar_menu.5.jsonnet

@@ -35,7 +35,8 @@ limitations under the License.
             ],
             garnish: null,
             served: "On The Rocks",
-        } for sd in [
+        }
+        for sd in [
             { name: "Yellow ", fruit: "Lemonade" },
             { name: "", fruit: "Orange Juice" },
         ]