scala pretty printer

Doc.scala

package org.im
package output
package prettyprinter

import scala.language._

import java.io.Writer
import java.time._

/**
 * Basic pretty printer object. Use combinators to combine instances
 * with others to build up output.
 *
 * From: http://scala-language.1934581.n4.nabble.com/file/n1995937/zPPrint.scala
 * Instructions: http://research.microsoft.com/en-us/um/people/daan/download/pprint/pprint-letter.pdf
 * Another lib: https://code.google.com/p/kiama/wiki/PrettyPrinting
 * Another paper: http://macbeth.cs.ucdavis.edu/ph-final.pdf
 * Another: https://speakerdeck.com/inkytonik/installing-trampolines-in-kiamas-pretty-printer
 * Another: https://code.google.com/p/kiama/source/browse/src/org/kiama/output/PrettyPrinter.scala?r=f9d6a7847f1d6b0bce5e6670c53a53f98b61e338
 */
sealed trait Doc {
  import Doc._

  /**
   * Convert to a string using ribbon 1.0 and width 80.
   */
  override def toString = {
    val writer = new java.io.StringWriter(2048)
    layout(writer, 1.0, 80, this).toString
  }

  /**
   * Convert to a string using the specified ribbon and width.
   */
  def show(ribbon_frac: Double, page_width: Int): String = {
    val wr = new java.io.StringWriter
    layout(wr, ribbon_frac, page_width, this).toString
  }

  /**
   * Write to a `Writer` with the specified ribbon and width.
   */
  def show(out: Writer, ribbon_frac: Double, page_width: Int): Writer =
    layout(out, ribbon_frac, page_width, this)

  /**
   * Compose two  documents. `lhs` goes on left side e.g. it is prepended.
   * Originally <>.
   */
  def ::(lhs: Doc): Doc = Cons(lhs, this)

  /**
   * Compose two documents together with a space between.
   * Originally <+>.
   */
  def :+:(lhs: Doc): Doc = lhs :: space :: this

  /**
   * Put this and lhs together horizontally with a space between them
   * if there is space on the current line or
   * lhs underneath this. Originally </>.
   */
  def :|:(lhs: Doc): Doc = lhs :: softline :: this

  /**
   * Put this and lhs directly together horizontally, directly concatenated,
   * or lhs underneath this. Originally <//>.
   */
  def :||:(lhs: Doc): Doc = lhs :: softbreak :: this

  /**
   * Compose with line. Put lhs below this. Originally <$>.
   */
  def :#:(lhs: Doc): Doc = lhs :: line :: this

  /**
   * Compose with `linebreak`. Put lhs below this.
   * Originally <$$>.
   */
  def :##:(lhs: Doc): Doc = lhs :: linebreak :: this

  /**
   * Compose two docs but dropping any empty documents, if any.
   * Like Lindig's ^^^.
   */
  def :%:(lhs: Doc): Doc = if (this == Empty) lhs
  else if (lhs == Empty) this
  else lhs :#: this
}

case object Empty extends Doc
case class Cons(h: Doc, t: Doc) extends Doc
case class Text(s: String) extends Doc
case class Nest(i: Int, d: Doc) extends Doc
case class Break(s: String) extends Doc
case class Group(d: Doc) extends Doc
case class Column(f: Int => Doc) extends Doc
case class Nesting(f: Int => Doc) extends Doc

// Mode (flat/break)
sealed trait Mode
case object FLT extends Mode
case object BRK extends Mode

/**
 * Module for `Doc`. Import to obtain unprefixed access to the
 * combinators. You can combine these with scala's string interpolation
 * to create strings e.g. formatted interpolator `f`.
 */
object Doc {

  /**
   * Typeclass for converting values to Doc objects.
   */
  trait ToDoc[T] {
    def doc(v: T): Doc
  }

  /**
   * Import converters to get automatic conversion from a value to a `Doc`.
   * Use `import Doc.Converters._`. Write your own and bring them
   * into scope to automatically convert values.
   */
  object Converters {
    implicit object StringToDoc extends ToDoc[String] {
      def doc(v: String) = string(v)
    }
    implicit object IntToDoc extends ToDoc[Int] {
      def doc(v: Int) = text(v toString)
    }
    implicit object LongToDoc extends ToDoc[Long] {
      def doc(v: Long) = text(v toString)
    }
    implicit object BooleanToDoc extends ToDoc[Boolean] {
      def doc(v: Boolean) = text(v toString)
    }
    implicit object ByteToDoc extends ToDoc[Byte] {
      def doc(v: Byte) = text(v toString)
    }
    implicit object ShortToDoc extends ToDoc[Short] {
      def doc(v: Short) = text(v toString)
    }
    implicit object CharToDoc extends ToDoc[Char] {
      def doc(v: Char) = text(v toString)
    }
    implicit object FloatToDoc extends ToDoc[Float] {
      def doc(v: Float) = text(v toString)
    }
    implicit object DoubleToDoc extends ToDoc[Double] {
      def doc(v: Double) = text(v toString)
    }
    implicit object LocalDateTimeToDoc extends ToDoc[LocalDateTime] {
      def doc(v: LocalDateTime) = text(v toString)
    }
    implicit object DocToDoc extends ToDoc[Doc] {
      def doc(v: Doc) = v
    }

    /**
     * Automatically lift values to a doc.
     */
    implicit def toDoc[T: ToDoc](v: T): Doc = implicitly[ToDoc[T]].doc(v)
  }

  /**
   * Import for syntax support.
   */
  object syntax {
    /**
     * Convert optional values to Doc objects or empty if the arg is None.
     * Use `optValue.doc`.
     *
     * TODO: Generalize to any container.
     */
    implicit class RichOption[T: ToDoc](s: Option[T]) {
      def doc: Doc = s.map(implicitly[ToDoc[T]].doc(_)) getOrElse empty
    }

    /**
     * Convert values to Doc objects. Use `value.doc`.
     */
    implicit class RichObject[T](v: T)(implicit todoc: ToDoc[T]) {
      def doc: Doc = todoc.doc(v)
    }
  }

  /**
   * Convenience formatter for a number using standard java semantics.
   * You can curry the function for easy re-use. Calls `String.format()`.
   */
  def f(format: String)(n: Number) = String.format(format, n)

  /**
   * Keep it dynamic to pick any changes in the property :-)
   */
  private[this] def lineTerminator = util.Properties.lineSeparator

  /**
   * Replaces all separators with `line`.
   */
  def string(s: String): Doc = s.split(lineTerminator).toList match {
    case Nil => empty
    case x :: xs => xs.foldLeft(text(x)) { _ :#: text(_) }
  }

  /**
   * Replaces all line separators with `softline`.
   */
  def softString(s: String): Doc = s.split(lineTerminator).toList match {
    case Nil => empty
    case x :: xs => xs.foldLeft(text(x)) { _ :|: text(_) }
  }

  /**
   * An empty document.
   */
  val empty: Doc = Empty

  /**
   * Convert a string to a document. If the string is empty, the
   * empty document is set.
   */
  def text(s: String): Doc = if (s == "") empty else Text(s)

  /**
   * Layout the doc on the current line then indent the next line and all subsequent docs.
   */
  def nest(i: Int, d: Doc): Doc = Nest(i, d)

  /**
   * Create a doc using the current column position to layout
   * the content dynamically.
   */
  def column = Column.apply _

  /**
   * Create a doc using the current nest column position for the line as an argument.
   * `nesting` can be anywhere on the line and not just at the initial
   * position of the nested line.
   */
  def nesting = Nesting.apply _

  /**
   * Keep doc on the same line by removing any line breaks if that line will
   * fit on the same line. Otherwise render the line breaks as is. The idea
   * is to group together docs and keep them together unless the grouped
   * docs are too large to fit.
   */
  def group(doc: Doc): Doc = Group(doc)

  /**
   * Advance to the next line. If undone by `group` then
   * concatenate with a space.
   */
  def line = Break(" ")

  /**
   * Advance to the next line. If undone by `group` then
   * concatenate directly together.
   */
  def linebreak = Break("")

  /**
   * Keep two docs horizontal and compose with a space between unless
   * there is not enough space on the current line.
   */
  def softline = group(line)

  /**
   * Keep two docs horizontal and compose with no space between them unless
   * there is not enough space on the current line.
   */
  def softbreak = group(linebreak)

  /**
   * Enclose a document with single quotes.
   */
  def squotes = enclose(squote, squote)_

  /**
   * Enclose a document in double quotes.
   */
  def dquotes = enclose(dquote, dquote)_

  /**
   * Enclose a document with braces "{" and "}".
   */
  def braces = enclose(lbrace, rbrace)_

  /**
   * Enclose a document with parenthesis "(" and ")"
   */
  def parens = enclose(lparen, rparen)_

  /**
   * Enclose a document in angled brackets.
   */
  def angles = enclose(langle, rangle)_

  /**
   * Encnlose a document in square brackets.
   */
  def brackets = enclose(lbracket, rbracket)_

  /**
   * Enclose with vertical bars.
   */
  def vbars = enclose(vert, vert)_

  /**
   * Only a left vertical bar.
   */
  def vbarsl(d: Doc) = vert :: d

  /**
   * Only a right vertical bar.
   */
  def vbarsr(d: Doc) = d :: vert

  /**
   * Enclose a document with `l` and `r` on the ends.
   */
  def enclose(l: Doc, r: Doc)(x: Doc) = l :: x :: r

  // symbols
  def lparen = text("(")
  def rparen = text(")")
  def langle = text("<")
  def rangle = text(">")
  def lbrace = text("{")
  def rbrace = text("}")
  def lbracket = text("[")
  def rbracket = text("]")
  def squote = text("\'")
  def dquote = text("\"")
  def semi = text(";")
  def colon = text(":")
  def comma = text(",")
  def space = text(" ")
  def dot = text(".")
  def backslash = text("\\")
  def assign = text("=")
  def vert = text("|")

  /**
   * Fold a function that takes two documenst over a list of documents.
   */
  def fold(op: (Doc, Doc) => Doc)(ld: Seq[Doc]): Doc = ld match {
    case Nil => empty
    case _ => ld reduceLeft { op(_: Doc, _: Doc) }
  }

  /**
   * After 'd', fill with spaces up to `fl` column unless the
   * width of doc is too wide in which case it increases the nesting
   * level to fl and inserts a line break.
   */
  def fillBreak(fl: Int, d: Doc): Doc =
    width(d, w => if (w > fl) nest(fl, linebreak) else text(spaces(fl - w)))

  /**
   * After `d`, fill with spaces up to `fl` column unless the width
   * of doc is too large in which case no spaces are appended.
   * Performs a right pad.
   * 
   * {{{
   * def stat[T](n: String, v: T) = fill(10, n) :+: ":" :+: v.toString
   * val params = Seq(
   *   ("Param1", s"${v1}"),
   *   ("My Param 2", s"${v2}"),
   *   ("Big Param 3", s"${v3}"))
   * println("Program parameters:")
   * println(align(vcat(params.map(t => stat(t._1, t._2)))).toString)
   * }}}
   * Produces:
   * {{{
   * Program Parameters:
   * Param1    : 1
   * My Param 2: 2
   * Big Param 3 : 3
   * }}}
   */
  def fill(fl: Int, d: => Doc): Doc =
    width(d, w => if (w >= fl) empty else text(spaces(fl - w)))

  /**
   * Given a doc and the column position after the doc, return a doc that will fill
   * the remaining space. This is a doc transducer.
   */
  def width(d: Doc, f: Int => Doc): Doc = column(k1 => d :: column(k2 => f(k2 - k1)))

  /**
   * Indent a document by `i` from the current position.
   */

  def indent(i: Int, d: Doc): Doc = hang(i, text(spaces(i)) :: d)

  /**
   * Create a hanging indentent from the current position. Subsequent lines
   * will display an indent.
   */
  def hang(i: Int, d: Doc): Doc = align(nest(i, d))

  /**
   * Render the document with the nesting level set to the current
   * column. It ignores the current nesting level. Subsequent lines
   * will show the alignment.
   */
  def align(d: Doc): Doc = column(k => nesting(i => nest(k - i, d)))

  /**
   * Concatenate all documents horizontally using :+: (docs separated
   * by a space) or vertically
   */
  def sep = group _ compose vsep

  /**
   * Fold over a list of docs with a soft line break between them.
   */
  def fillSep = fold(_ :|: _) _

  /**
   * Fold over a list of docs with a space between each.
   */
  def hsep = fold(_ :+: _) _

  /**
   * Fold over a list of docs with a line between them.
   */
  def vsep = fold(_ :#: _) _

  /**
   * Group together the doc to try and keep them horizontal, if it fits the page,
   * or vertically if not.
   */
  def cat = group _ compose vcat

  /**
   * Concatenate horizontally as long as it fits the page then
   * insert a `softbreak` and continue concatinating.
   */
  def fillCat = fold(_ :||: _) _

  /**
   * Fold over a list of docs without any separators (::).
   */
  def hcat = fold(_ :: _) _

  /**
   * Fold over a list of docs with a break (:##:) between each item.
   */
  def vcat = fold(_ :##: _) _

  /**
   * Limit the width of string. Show marker if text is truncated. Otherwise
   * pad to the requested length. If there is not enough space for marker,
   * just return a dot ".".
   *
   * TODO: Not sure this works with any marker.
   */
  def limit(width: Int, marker: String = "...")(content: => String): Doc = {
    require(width > 0)
    val len = content.size
    if (len == 0) spaces(width)
    if (len < width) string(content + spaces(width - len))
    else if (len > width && width > 3) string(content.substring(0, width - marker.size) + marker)
    else if (len > width && width <= 1) text(".")
    else if (len > width && width <= 2) string(content.charAt(0) + ".")
    else string(content)
  }

  /**
   * Concatenates all documents with `p` except for the last document.
   */
  def punctuate(p: Doc, docs: Seq[Doc]): List[Doc] = docs match {
    case Nil => Nil
    case one@d :: Nil => one
    case d :: ds => (d :: p) :: punctuate(p, ds)
  }

  /**
   * Print using a list look and feel. List means
   * square bracket: `[doc,` which is more haskell/python
   * than scala.
   */
  def list = encloseSep(lbracket, rbracket, comma)_

  /**
   * Print using tuple look and feel e.g. `(doc,`.
   */
  def tupled = encloseSep(lparen, rparen, comma)_

  /**
   * Print using brace look and feel e.g. `{doc,`.
   */
  def semiBraces = encloseSep(lbrace, rbrace, semi)_

  /**
   * Enclose a doc with a left, right and separator doc. If the doc goes vertical,
   * align subsequent lines on the first entry.
   */
  def encloseSep(l: Doc, r: Doc, sep: Doc)(ds: Seq[Doc]): Doc = ds match {
    case Nil => l :: r
    case x :: Nil => l :: x :: r
    case x :: xs => align(cat((l :: x) :: xs.map(sep :: _)) :: r)
  }

  /**
   * Layout the text in plain text. Without tail recursion, this
   * could blow the stack.
   *
   * @param out Writer to emit content to.
   * @param frac ratio of ribbon to full page width
   * @param w full page width
   * @return The input Writer
   */
  def layout(out: Writer, frac: Double, w: Int, doc: Doc): Writer = {

    type Cells = Seq[(Int, Mode, Doc)]

    val ribbon = (frac * w round).asInstanceOf[Int] min w max 0

    /**
     * Runs forward to see if a doc fits.
     *
     * w - ribbon space left
     * k - current column
     */
    def fits(w: Int, k: Int, cs: Cells): Boolean = cs match {
      case _ if w < 0 => false
      case Nil => true
      case (i, _, Empty) :: z => fits(w, k, z)
      case (i, m, Cons(x, y)) :: z => fits(w, k, (i, m, x) :: (i, m, y) :: z)
      case (i, m, Nest(j, x)) :: z => fits(w, k, (i + j, m, x) :: z)
      case (i, m, Text(s)) :: z => fits(w - s.length, k + s.length, z)
      case (i, FLT, Break(s)) :: z => fits(w - s.length, k + s.length, z)
      case (i, BRK, Break(_)) :: z => true
      case (i, m, Group(x)) :: z => fits(w, k, (i, m, x) :: z)
      case (i, m, Column(f)) :: z => fits(w, k, (i, m, f(k)) :: z)
      case (i, m, Nesting(f)) :: z => fits(w, k, (i, m, f(i)) :: z)
    }

    /**
     * Emit a string to a `Writer` and return the `Writer`.
     */
    def emit(w: Writer, s: String): Writer = { w write s; w }

    /**
     * Emit a new newline then add `i` spaces. Return the `Writer`.
     */
    def nl(acc: Writer, i: Int): Writer = spaces(emit(acc, "\n"), i)

    // best
    // n - indentation of current line
    // k - current column
    def best(acc: Writer, n: Int, k: Int, cs: Cells): Writer = cs match {
      case Nil => acc
      case (i, _, Empty) :: z => best(acc, n, k, z)
      case (i, m, Cons(x, y)) :: z => best(acc, n, k, (i, m, x) :: (i, m, y) :: z)
      case (i, m, Nest(j, x)) :: z => best(acc, n, k, (i + j, m, x) :: z)
      case (i, _, Text(s)) :: z => best(emit(acc, s), n, k + s.length, z)
      case (i, FLT, Break(s)) :: z => best(emit(acc, s), n, k + s.length, z)
      case (i, BRK, Break(s)) :: z => best(nl(acc, i), i, i, z)
      case (i, FLT, Group(x)) :: z => best(acc, w, k, (i, FLT, x) :: z)
      case (i, BRK, Group(x)) :: z => {
        val ribbonleft = (w - k) min (ribbon - k + n)
        if (fits(ribbonleft, k, (i, FLT, x) :: z))
          best(acc, n, k, (i, FLT, x) :: z)
        else
          best(acc, n, k, (i, BRK, x) :: z)
      }
      case (i, m, Column(f)) :: z => best(acc, n, k, (i, m, f(k)) :: z)
      case (i, m, Nesting(f)) :: z => best(acc, n, k, (i, m, f(i)) :: z)
    }
    best(out, 0, 0, (0, BRK, doc) :: Nil)
  }

  /**
   * Write out spaces to `acc`.
   */
  private def spaces(acc: Writer, n: Int): Writer = {
    var c = ' '
    var i = n
    while (i > 0) { acc append c; i -= 1 }
    acc
  }

  /**
   * Create a string of spaces, or some other char, with the specified length.
   */
  def spaces(n: Int, c: Character = ' '): String = {
    val acc = new StringBuffer(80)
    var i = n
    while (i > 0) { acc append c; i -= 1 }
    acc toString
  }
  
  /**
   * For each double quote in `in`, double up the quote. This effectively escapes
   * it for most CSV readers.
   */
  def escapedquote(in: String): String = in.replace("\"", "\"\"")
}