A Scala 3 port of the code of Document.scala from the Scala 2 stdlib

Document.scala

/*
  Scala 3 port of Michel Schinz's Document.scala
  A basic pretty-printing library, based on Lindig's strict version
  of Wadler's adaptation of Hughes' pretty-printer.

  https://github.com/scala/scala/blob/v2.11.8/src/library/scala/text/Document.scala

  Remove the example and Main at the end if you want to use this in your code.

  Example of usage with an AST can be run directly with
  > scala Document.scala

  Author: M Nimalan (@mark1626)

  This is free and unencumbered software released into the public domain.
**/
package com.mark

import java.io.Writer

object PrettyPrint:
  enum Doc:
    case DocNil
    case DocBreak
    case DocSpace
    case DocText(txt: String)
    case DocCons(hd: Doc, tl: Doc)
    case DocNest(indent: Int, doc: Doc)

    def <>(hd:Doc): Doc = (this, hd) match
      case (_, DocNil) => this
      case (DocNil, _) => hd
      case _ => DocCons(this, hd)

    def <@>(hd: Doc): Doc = hd match
      case DocNil => this
      case _ => this <> DocBreak <> hd

    def <+>(hd: Doc): Doc = this <> DocSpace <> hd

    def pretty: String =
      val writer = new java.io.StringWriter()
      format(writer)
      writer.toString

    def format(writer: Writer): Unit =
      type FmtState = (Int, Doc)

      def spaces(n: Int): Unit =
        var rem = n
        while (rem >= 16) { writer write "                "; rem -= 16 }
        if (rem >= 8) { writer write "        "; rem -= 8 }
        if (rem >= 4) { writer write "    "; rem -= 4 }
        if (rem >= 2) { writer write "  "; rem -= 2 }
        if (rem >= 1) { writer write " " }

      def fmt(state: List[FmtState]): Unit = state match
        case List() => ()
        case (_, DocNil) :: z => fmt(z)
        case (i, DocCons(h, t)) :: z => fmt((i, h) :: (i, t) :: z)
        case (_, DocText(t)) :: z =>
          writer.write(t)
          fmt(z)
        case (i, DocNest(ii, d)) :: z => fmt((i+ii, d):: z)
        case (i, DocBreak) :: z =>
          writer.write("\n")
          spaces(i)
          fmt(z)
        case (_, DocSpace) :: z =>
          writer.write(" ")
          fmt(z)
        case _ => ()

      fmt(List((0, this)))

  object Doc:
    def emptyDoc = DocNil
    def text(s: String): Doc = DocText(s)
    def value(v: Any): Doc = DocText(v.toString)
    def space = DocSpace
    def line = DocBreak

    def semi: Doc = text(";")
    def colon: Doc = text(":")
    def comma: Doc = text(",")
    def dot: Doc = text(".")
    def equal: Doc = text("=")
    def lbrace: Doc = text("{")
    def rbrace: Doc = text("}")

    def nest(d: Doc, i: Int): Doc = DocNest(i, d)

    def folddoc(ds: Iterable[Doc], f: (Doc, Doc) => Doc) = ds.isEmpty match
      case true => emptyDoc
      case false => ds.tail.foldLeft(ds.head)(f)

    def hsep(ds: Iterable[Doc]): Doc =
      folddoc(ds, _ <+> _)

    def vsep(ds: Iterable[Doc]): Doc =
      folddoc(ds, _ <@> _)

    def enclose(l: Doc, d: Doc, r: Doc): Doc = hsep(Seq(l, d, r))
    def parens(d: Doc) = enclose(text("("), d, text(")"))
    def braces(d: Doc) = enclose(lbrace, d, rbrace)

    def scope(
      doc: Doc,
      left: Doc = lbrace,
      right: Doc = rbrace,
      indent: Int = 2
    ): Doc =
      left <> nest(emptyDoc <@> doc, indent) <@> right

/*
  A toy AST definition to test the pretty printer
*/

import PrettyPrint.Doc
import PrettyPrint.Doc._

enum BOp:
  case Add
  case GT

enum Expr:
  case EVar(id: String)
  case EInt(v: Int)
  case EBinop(op: BOp, e1: Expr, e2: Expr)

enum Statement:
  case SAssign(lhs: Expr, rhs: Expr)
  case SPar(stmts: Seq[Statement])
  case SIf(cond: Expr, cons: Statement, alt: Statement)
  case SEmpty

// Extensions to the ADT to pretty print

import BOp.*
import Statement.*
import Expr.*

extension (op: BOp)
  def emit: Doc = op match
    case Add => text("+")
    case GT => text(">")

extension (expr: Expr)
  def emit: Doc = expr match
    case EVar(id) => text(id)
    case EInt(v) => text(v.toString)
    case EBinop(op, e1, e2) => e1.emit <+> op.emit <+> e2.emit

extension (stmt: Statement)
  def emit: Doc = stmt match
    case SAssign(lhs, rhs) => lhs.emit <+> equal <+> rhs.emit <> semi
    case SPar(stmts) => scope(vsep(stmts.map(stmt => stmt.emit)))
    case SIf(cond, cons, alt) => text("if") <+> parens(cond.emit) <@> cons.emit
    case SEmpty => text("")


/*
{
  x = 5;
  z = 35;
  a = 25;
  {
    y = 10;
    x = x + y;
    if ( x > 10 )
    {
      y = 5;
    }
  }
}
*/
val prog: Statement = SPar(Seq(
  SAssign(EVar("x"), EInt(5)),
  SAssign(EVar("z"), EInt(35)),
  SAssign(EVar("a"), EInt(25)),
  SPar(Seq(
    SAssign(EVar("y"), EInt(10)),
    SAssign(EVar("x"), EBinop(BOp.Add, EVar("x"), EVar("y"))),
    SIf(
      EBinop(BOp.GT, EVar("x"), EInt(10)),
      SPar(Seq(
        SAssign(EVar("y"), EInt(5))
      )),
      SEmpty)
  ))
))

@main def main: Unit =
  println(prog.emit.pretty)