implementation of the Wadler document algebra for strict languages described by Lindig (2000)

gistfile1.ex

defmodule WadlerLindig do
  # Elixir implementation of the Wadler document algebra as described in 
  # "Strictly Pretty" (2000) by Christian Lindig
  # http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.34.2200
  # 
  # The original Haskell implementation of the algorithm relies on lazy 
  # evaluation to unfold document groups on two alternatives:
  # _flat_ (breaks as spaces) and _broken_ (breakes as newlines).
  # Implementing the same logic on a strict language such as Elixir leads
  # to an exponential growth of the possible documents, unless document
  # groups are encoded explictly as _flat_ or _broken_. Those groups are
  # then reduced to a simple document, where the layout is already decided.

  # shortcuts (appendix A: missing parts)
  def newline, do: "\n"
  def strlen(s), do: String.length(s)

  def repeat(_, 0), do: ""
  def repeat(s, i), do: String.duplicate s, i

  # document type
  @type doc :: DocNil | DocCons.t | DocText.t | DocNest.t | DocBreak.t | DocGroup.t
  defrecord DocCons, left: DocNil, right: DocNil
  defrecord DocText, str: ""
  defrecord DocNest, indent: 1, doc: DocNil
  defrecord DocBreak, str: " "
  defrecord DocGroup, doc: DocNil

  @spec concat(doc, doc) :: DocCons.t
  def concat(x, y), do: DocCons[left: x, right: y]
  
  @spec empty() :: DocNil
  def empty, do: DocNil

  @spec text(binary) :: DocText.t
  def text(s) when is_binary(s), do: DocText[str: s]

  @spec nest(non_neg_integer, doc) :: DocNest.t
  def nest(i, x) when is_integer(i), do: DocNest[indent: i, doc: x]

  @spec break(binary) :: DocBreak.t
  def break(s) when is_binary(s), do: DocBreak[str: s]

  @spec break() :: DocBreak.t
  def break(), do: DocBreak[str: " "]

  @spec group(doc) :: DocGroup.t
  def group(d), do: DocGroup[doc: d]

  # simple document type
  @type sdoc :: SNil | SText.t | SLine.t
  defrecord SText, str: "", sdoc: SNil
  defrecord SLine, indent: 1, sdoc: SNil # newline + spaces

  @spec render(sdoc) :: binary
  def render(SNil), do: ""
  def render(SText[str: s, sdoc: d]), do: s <> render(d)
  def render(SLine[indent: i, sdoc: d]) do
    prefix = repeat " ", i
    newline <> prefix <> render d
  end

  # document mode: __flat__ or __broken__
  @type mode :: Flat | Break

  # the fits? and format methods have to deal explicitly with the document modes
  @spec fits?(integer, { integer, mode, doc }[]) :: boolean
  def fits?(w, _) when w < 0,                              do: false
  def fits?(_, []),                                        do: true
  def fits?(w, [{_, _, DocNil} | t]),                      do: fits? w, t
  def fits?(w, [{i, m, DocCons[left: x, right: y]} | t]),  do: fits? w, [{i, m, x} | [{i, m, y} | t]]
  def fits?(w, [{i, m, DocNest[indent: j, doc: x]} | t]),  do: fits? w, [{i + j, m, x} | t]
  def fits?(w, [{_, _, DocText[str: s]} | t]),             do: fits? (w - strlen s), t
  def fits?(w, [{_, Flat, DocBreak[str: s]} | t]),         do: fits? (w - strlen s), t
  def fits?(_, [{_, Break, DocBreak[str: _]} | _]),        do: true # impossible (but why?)
  def fits?(w, [{i, _, DocGroup[doc: x]} | t]),            do: fits? w, [{i, Flat, x} | t]

  @spec format(integer, integer, { integer, mode, doc }[]) :: boolean
  def format(_, _, []),                                       do: SNil
  def format(w, k, [{_, _, DocNil} | t]),                     do: format w, k, t
  def format(w, k, [{i, m, DocCons[left: x, right: y]} | t]), do: format w, k, [{i, m, x} | [{i, m, y} | t]]
  def format(w, k, [{i, m, DocNest[indent: j, doc: x]} | t]), do: format w, k, [{i + j, m, x} | t]
  def format(w, k, [{_, _, DocText[str: s]} | t]),            do: SText[str: s, sdoc: format w, (k + strlen s), t]
  def format(w, k, [{_, Flat, DocBreak[str: s]} | t]),        do: SText[str: s, sdoc: format w, (k + strlen s), t]
  def format(w, _, [{i, Break, DocBreak[str: _]} | t]),       do: SLine[indent: i, sdoc: format w, i, t]
  def format(w, k, [{i, _, DocGroup[doc: x]} | t]) do
    if fits? (w - k), [{i, Flat, x} | t] do
      format w, k, [{i, Flat, x} | t]
    else
      format w, k, [{i, Break, x} | t]
    end
  end

  # interface (appendix A)
  @spec pretty(integer, doc) :: binary
  def pretty(w, d) do
    sdoc = format w, 0, [{0, Flat, DocGroup[doc: d]}]
    render(sdoc)
  end

  # functions implementing the if-then-else example on section 4
  def opt_break(x, y) do 
    case {x, y} do
      {DocNil, _} -> y
      {_, DocNil} -> x
      _           -> concat(x, concat(break, y))
    end
  end

  def binop(left, op, right) do
    group(nest(2, 
      group(
        opt_break(
          opt_break(text(left), text(op)),
          text(right)
        )
      )
    ))
  end

  def if_then_else(c, d, e) do
    group(
      opt_break(opt_break(
          group(nest(2, opt_break(text("if"),   c))),
          group(nest(2, opt_break(text("do"),   d)))),
          group(nest(2, opt_break(text("else"), e)))
      )
    )
  end

  def example do
    c = binop("a", "==", "b")
    d = binop("a", "<<", "2")
    e = binop("a", "+",  "b")

    doc = if_then_else(c, d, e)

    Enum.map [32,15,8,7,6], fn(n) ->
      title = "w = #{n}"
      dashes = repeat "-", strlen(title)
      IO.puts "\n#{title}\n#{dashes}"
      IO.puts pretty(n, doc)
    end
  end
end

WadlerLindig.example