Doc.sk

/**
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

// Implements the intermediate representation described in
// "A prettier printer" by Philip Wadler
// http://homepages.inf.ed.ac.uk/wadler/papers/prettier/prettier.pdf
// A more friendly explanation is available on the prettier docs:
// https://github.com/prettier/prettier/blob/master/commands.md

module Doc;
const space: Doc = Doc.Str(" ");

base class .Doc() {}

class Empty() extends Doc {}
class Concat(parts: Vector<Doc>) extends Doc {
  static fun createFromItems<T: Sequence<Doc>>(elements: T): Concat {
    Concat(elements.collect(Vector))
  }
  fun inspect(): Inspect {
    InspectVector("Doc.Concat", this.parts.map(inspect).toArray());
  }
}
class Indent(contents: Doc) extends Doc {
  static fun createFromItems<T: Sequence<Doc>>(elements: T): Indent {
    Indent(Concat(elements.collect(Vector)))
  }
  fun inspect(): Inspect {
    this.contents match {
    | Concat(parts) -> InspectVector("Doc.Indent", parts.map(inspect).toArray())
    | contents -> InspectCall("Doc.Indent", Array[inspect(contents)])
    }
  }
}
class Align(n: Int, contents: Doc) extends Doc {}
class Group(
  contents: Doc,
  breakLine: Bool = false,
  marker: ?Marker = None(),
) extends Doc {
  static fun createFromItems<T: Sequence<Doc>>(elements: T): Group {
    Group(Concat(elements.collect(Vector)))
  }
  fun inspect(): Inspect {
    this match {
    | Group(Concat(parts), false, None()) ->
      InspectVector("Doc.Group", parts.map(inspect).toArray())
    | _ ->
      InspectCall(
        "Doc.Group",
        Array[
          inspect(this.contents),
          inspect(this.breakLine),
          inspect(this.marker),
        ],
      )
    }
  }
}

// The default grouping algorithm is to check if the entire content would fit,
// otherwise break. This works for most cases except for one: inlining lambda:
//   call(a, b, x -> {
//   })
// What we want is to ensure that 'a, b, x -> {' fits in one line, otherwise
// break:
//   call(
//     a,
//     b,
//     x -> {
//     }
//   )
// In order to do that, we introduce the notion of a marker. The group is
// inlined if from the beginning of the group to the first marker fits in one
// line.
class Marker(id: TextRange) uses Equality extends Doc {}

class IfBreak(breakContents: Doc, flatContents: Doc = Empty()) extends Doc {}

class Str(str: String) extends Doc {}

base class GenericLine() extends Doc {}
class Line() extends GenericLine {}
class SoftLine() extends GenericLine {}
class HardLine() extends GenericLine {}

class LineSuffix(contents: Doc) extends Doc {}
class LineSuffixBoundary() extends Doc {}

// Fills each line with as much code as possible before moving to a new
// line with the same indentation.
//
// Expects doc.parts to be a Vector of alternating content and
// whitespace. The whitespace contains the linebreaks.
//
// For example:
//   ["I", line, "love", line, "monkeys"]
// or
//   [{ type: group, ... }, softline, { type: group, ... }]
//
// It uses this parts structure to handle three main layout cases:
// * The first two content items fit on the same line without
//   breaking
//   -> output the first content item and the whitespace "flat".
// * Only the first content item fits on the line without breaking
//   -> output the first content item "flat" and the whitespace with
//   "breakLine".
// * Neither content item fits on the line without breaking
//   -> output the first content item and the whitespace with "breakLine".
class Fill(parts: Vector<Doc>) extends Doc {}

private base class Mode {
  children =
  | Break()
  | Flat()
}
private class IndentComputation(
  indent: Int,
  alignSpaces: Int,
  alignTabs: Int,
) {}
private class Command(indent: IndentComputation, mode: Mode, doc: Doc) {}

private fun rootIndent(): IndentComputation {
  IndentComputation(0, 0, 0)
}

private fun makeIndent(ind: IndentComputation): IndentComputation {
  IndentComputation(ind.indent + 1, ind.alignSpaces, ind.alignTabs)
}

private fun makeAlign(ind: IndentComputation, n: Int): IndentComputation {
  if (n == -Int::max) {
    IndentComputation(0, 0, 0)
  } else {
    IndentComputation(
      ind.indent,
      ind.alignSpaces + n,
      ind.alignTabs + (if (n > 0) 1 else 0),
    )
  }
}

private fun pop<T>(buffer: mutable Vector<T>): ?T {
  if (buffer.size() == 0) {
    None()
  } else {
    Some(buffer.pop())
  }
}

private fun fits(
  next: Command,
  restCommands: mutable Vector<Command>,
  width: Int,
  marker: ?Marker = None(),
): Bool {
  restIdx = restCommands.size();
  cmds = mutable Vector[next];
  while (width >= 0) {
    pop(cmds) match {
    | None() ->
      if (restIdx == 0) {
        return !marker.isSome()
      };
      cmds.push(restCommands[restIdx - 1]);

      !restIdx = restIdx - 1
    | Some(Command(ind, mode, doc)) ->
      doc match {
      | Str(str) -> !width = width - str.length()
      | Concat(parts) ->
        parts.eachWithIndex((i, _) -> {
          cmds.push(Command(ind, mode, parts[parts.size() - i - 1]))
        })
      | Indent(contents) -> cmds.push(Command(makeIndent(ind), mode, contents))
      | Align(n, contents) ->
        cmds.push(Command(makeAlign(ind, n), mode, contents))
      | Group(contents, breakLine, _marker) ->
        // Note(vjeux): I don't know if it's safe to ignore breaks when
        // there's a marker. The group that breakLine are above the one with
        // the parent marker.
        newMode = if (breakLine && !marker.isSome()) Break() else mode;
        cmds.push(Command(ind, newMode, contents))
      | IfBreak(breakContents, flatContents) ->
        mode match {
        | Break() -> cmds.push(Command(ind, mode, breakContents))
        | Flat() -> cmds.push(Command(ind, mode, flatContents))
        }
      | line @ GenericLine() ->
        mode match {
        | Flat() ->
          line match {
          | Line() -> !width = width - 1
          | HardLine() -> return !marker.isSome()
          | SoftLine() -> void
          }
        | Break() -> return !marker.isSome()
        }

      | m @ Marker(_) ->
        if (marker == Some(m)) {
          return true
        }

      | Fill(parts) ->
        parts.reversed().each(part -> {
          cmds.push(Command(ind, mode, part))
        })
      | LineSuffix _
      | LineSuffixBoundary _
      | Empty() ->
        void
      }
    }
  };

  false
}

fun printDocToString(
  doc: Doc,
  width: Int = 80,
  tabWidth: Int = 2,
  useTabs: Bool = false,
  newLine: String = "\n",
): String {
  out = mutable Vector[];
  printDoc(doc, out.push, width, tabWidth, useTabs, newLine);
  out.join("")
}

fun printDoc(
  doc: Doc,
  print: String -> void,
  width: Int = 80,
  tabWidth: Int = 2,
  useTabs: Bool = false,
  newLine: String = "\n",
): void {
  (!doc, _) = propagateBreaks(doc);
  pos = 0;
  // cmds is basically a stack. We've turned a recursive call into a
  // while loop which is much faster. The while loop below adds new
  // cmds to the array instead of recursively calling `print`.
  cmds = mutable Vector[Command(rootIndent(), Break(), doc)];
  shouldRemeasure = false;

  out = mutable Vector[];
  lineSuffix = mutable Vector[];

  while ({
    pop(cmds) match {
    | None() ->
      if (lineSuffix.isEmpty()) {
        false
      } else {
        lineSuffix.reversed().each(cmd -> cmds.push(cmd));
        lineSuffix.clear();
        true
      }
    | Some(cmd @ Command(ind, mode, cmd_doc)) ->
      cmd_doc match {
      | Str(str) ->
        out.push(str);
        !pos = pos + str.length()
      | Concat(parts) ->
        parts.eachWithIndex((i, _) -> {
          cmds.push(Command(ind, mode, parts[parts.size() - i - 1]))
        })
      | Indent(contents) -> cmds.push(Command(makeIndent(ind), mode, contents))
      | Align(n, contents) ->
        cmds.push(Command(makeAlign(ind, n), mode, contents))
      | Group(contents, breakLine, marker) ->
        if (!marker.isSome() && mode is Flat() && !shouldRemeasure) {
          cmds.push(Command(ind, if (breakLine) Break() else Flat(), contents))
        } else {
          !shouldRemeasure = false;

          next = Command(ind, Flat(), contents);
          rem = width - pos;

          doesFit = fits(next, cmds, rem, marker);
          if (marker.isSome() && doesFit) {
            newContents = flattenUntilMarker(contents, marker.fromSome()).i0;
            cmds.push(Command(ind, Flat(), newContents));
            !shouldRemeasure = true
          } else if (!breakLine && doesFit) {
            cmds.push(next)
          } else {
            cmds.push(Command(ind, Break(), contents))
          }
        }
      | IfBreak(breakContents, flatContents) ->
        mode match {
        | Break() -> cmds.push(Command(ind, mode, breakContents))
        | Flat() -> cmds.push(Command(ind, mode, flatContents))
        }
      | LineSuffix(contents) -> lineSuffix.push(Command(ind, mode, contents))
      | LineSuffixBoundary() ->
        if (!lineSuffix.isEmpty()) {
          cmds.push(Command(ind, mode, Doc.HardLine()))
        }
      | line @ GenericLine() ->
        if (mode is Flat() && line is Line()) {
          out.push(" ");
          !pos = pos + 1
        } else if (mode is Flat() && line is SoftLine()) {
          void
        } else {
          if (mode is Flat() && line is HardLine()) {
            // This line was forced into the output even if we
            // were in flattened mode, so we need to tell the next
            // group that no matter what, it needs to remeasure
            // because the previous measurement didn't accurately
            // capture the entire expression (this is necessary
            // for nested groups)
            !shouldRemeasure = true
          };

          if (!lineSuffix.isEmpty()) {
            cmds.push(cmd);
            lineSuffix.reversed().each(cmd -> cmds.push(cmd));
            lineSuffix.clear()
          } else {
            length = ind.indent * tabWidth + ind.alignSpaces;
            indentString = if (useTabs) {
              "\t".repeat(ind.indent + ind.alignTabs)
            } else {
              " ".repeat(length)
            };
            if (out.size() > 0) {
              print(out.join("").trimRight());
              out.clear()
            };
            print(newLine);
            out.push(indentString);
            !pos = length
          }
        }
      | Fill(parts) ->
        rem = width - pos;

        if (parts.size() != 0) {
          content = parts[0];
          contentFlatCmd = Command(ind, Flat(), content);
          contentBreakCmd = Command(ind, Break(), content);
          contentFits = fits(contentFlatCmd, mutable Vector[], rem);

          if (parts.size() == 1) {
            if (contentFits) {
              cmds.push(contentFlatCmd)
            } else {
              cmds.push(contentBreakCmd)
            }
          } else {
            whitespace = parts[1];
            whitespaceFlatCmd = Command(ind, Flat(), whitespace);
            whitespaceBreakCmd = Command(ind, Break(), whitespace);

            if (parts.size() == 2) {
              if (contentFits) {
                cmds.push(whitespaceFlatCmd);
                cmds.push(contentFlatCmd)
              } else {
                cmds.push(whitespaceBreakCmd);
                cmds.push(contentBreakCmd)
              }
            } else {
              // At this point we've handled the first pair (context, separator)
              // and will create a new fill doc for the rest of the content.
              remainingCmd = Command(ind, mode, Fill(parts.slice(2)));

              secondContent = parts[2];

              firstAndSecondContentFlatCmd = Command(
                ind,
                Flat(),
                Concat[content, whitespace, secondContent],
              );
              firstAndSecondContentFits = fits(
                firstAndSecondContentFlatCmd,
                mutable Vector[],
                rem,
              );

              if (firstAndSecondContentFits) {
                cmds.push(remainingCmd);
                cmds.push(whitespaceFlatCmd);
                cmds.push(contentFlatCmd)
              } else if (contentFits) {
                cmds.push(remainingCmd);
                cmds.push(whitespaceBreakCmd);
                cmds.push(contentFlatCmd)
              } else {
                cmds.push(remainingCmd);
                cmds.push(whitespaceBreakCmd);
                cmds.push(contentBreakCmd)
              }
            }
          }
        }

      | Marker _
      | Empty() ->
        void
      };
      true
    }
  }) void;
  if (out.size() > 0) {
    print(out.join("").trimRight())
  }
}

private fun flattenUntilMarker(doc: Doc, marker: Doc.Marker): (Doc, Bool) {
  doc match {
  | Concat(parts) ->
    found = false;
    newParts = parts.map(part -> {
      if (!found) {
        (newPart, foundPart) = flattenUntilMarker(part, marker);
        if (foundPart) {
          !found = true
        };
        newPart
      } else {
        part
      }
    });
    (Concat(newParts), found)

  | Indent(contents) -> flattenUntilMarker(contents, marker)
  | Align(_, contents) -> flattenUntilMarker(contents, marker)
  | IfBreak(_breakContents, flatContents) ->
    flattenUntilMarker(flatContents, marker)

  | Group(contents, _break, groupMarker) ->
    (newContents, found) = flattenUntilMarker(contents, marker);
    if (!found) {
      (newContents, found)
    } else {
      (!newContents, newBreak) = propagateBreaks(newContents);
      (Group(newContents, newBreak, groupMarker), found)
    }

  | HardLine()
  | SoftLine() ->
    (Doc.Empty(), false)

  | Line() -> (Doc.space, false)

  | Fill _ // Need to implement it
  | LineSuffix _
  | LineSuffixBoundary()
  | Empty()
  | Str _ ->
    (doc, false)

  | m @ Marker _ -> (m, m == marker)
  }
}

private fun propagateBreaks(doc: Doc): (Doc, Bool) {
  doc match {
  | Concat(parts) ->
    shouldBreakParts = false;
    newParts = parts.map(part -> {
      (newPart, shouldBreak) = propagateBreaks(part);
      !shouldBreakParts = shouldBreakParts || shouldBreak;
      newPart
    });
    (Concat(newParts), shouldBreakParts)

  | Indent(contents) ->
    (newContents, shouldBreak) = propagateBreaks(contents);
    (Indent(newContents), shouldBreak)

  | Align(n, contents) ->
    (newContents, shouldBreak) = propagateBreaks(contents);
    (Align(n, newContents), shouldBreak)

  | Group(contents, breakLine, marker) ->
    (newContents, shouldBreak) = propagateBreaks(contents);
    groupShouldBreak = if (marker.isSome()) {
      false
    } else {
      breakLine || shouldBreak
    };
    (Group(newContents, groupShouldBreak, marker), breakLine || shouldBreak)

  | IfBreak(breakContents, flatContents) ->
    (newBreak, _) = propagateBreaks(breakContents);
    (newFlat, shouldBreak) = propagateBreaks(flatContents);
    (IfBreak(newBreak, newFlat), shouldBreak)

  | LineSuffix _
  | HardLine() ->
    (doc, true)

  | Fill _
  | Marker _
  | LineSuffixBoundary()
  | SoftLine()
  | Line()
  | Empty()
  | Str(_str) ->
    (doc, false)
  }
}

fun simplifyDoc(doc: Doc): Doc {
  doc match {
  // Flatten all the nested Concat()
  | Concat(parts) ->
    newParts = mutable Vector<Doc>[];

    pushPart = (_part: Doc) -> void;
    !pushPart = part -> {
      simplifyDoc(part) match {
      // Remove all the Empty()
      | Empty() -> void
      // Combine consecutive Str()
      | x @ Str(str) ->
        if (newParts.size() > 0) {
          newParts[newParts.size() - 1] match {
          | Str(lastStr) -> newParts![newParts.size() - 1] = Str(lastStr + str)
          | _ -> newParts.push(x)
          }
        } else {
          newParts.push(x)
        }
      | Concat(xParts) -> xParts.each(pushPart)
      | x -> newParts.push(x)
      }
    };

    parts.each(part -> pushPart(part));
    if (newParts.size() == 0) {
      Empty()
    } else if (newParts.size() == 1) {
      newParts[0]
    } else {
      Concat(newParts.chill())
    }

  | Indent(contents) ->
    simplifyDoc(contents) match {
    | Empty() -> Empty()
    | x -> Indent(x)
    }
  | Align(n, contents) ->
    simplifyDoc(contents) match {
    | Empty() -> Empty()
    | x -> Align(n, x)
    }
  | Group(contents, breakLine, marker) ->
    simplifyDoc(contents) match {
    | Empty() -> Empty()
    | x -> Group(x, breakLine, marker)
    }
  | IfBreak(breakContents, flatContents) ->
    (simplifyDoc(breakContents), simplifyDoc(flatContents)) match {
    | (Empty(), Empty()) -> Empty()
    | (x, y) -> IfBreak(x, y)
    }
  | Fill(parts) -> Fill(parts.map(simplifyDoc))

  | Str("") -> Empty()

  | Empty()
  | LineSuffix _
  | HardLine()
  | Marker _
  | LineSuffixBoundary()
  | SoftLine()
  | Line()
  | Str _ ->
    doc
  }
}

module end;