operators.scala

package fastparse

// Ported from http://hackage.haskell.org/package/megaparsec-6.5.0/docs/src/Text.Megaparsec.Expr.html#Operator
package object operators {

  import fastparse.all._

  sealed abstract class Operator[A]

  /** Non-associative infix */
  final case class InfixN[A](run: P[A => A => A]) extends Operator[A]

  /** Left-associative infix */
  final case class InfixL[A](run: P[A => A => A]) extends Operator[A]

  /** Right-associative infix */
  final case class InfixR[A](run: P[A => A => A]) extends Operator[A]

  /** Prefix */
  final case class Prefix[A](run: P[A => A]) extends Operator[A]

  /** Postfix */
  final case class Postfix[A](run: P[A => A]) extends Operator[A]

  private def const[A, B](x: A) = (y: B) => x

  def prefix[A](name: String)(f: A => A): Operator[A]      = Prefix(P(name).map(const(f)))
  def postfix[A](name: String)(f: A => A): Operator[A]     = Postfix(P(name).map(const(f)))
  def binary[A](name: String)(f: (A, A) => A): Operator[A] = InfixL(P(name).map(const(f.curried)))

  def makeExprParser[A](term: P[A], table: List[List[Operator[A]]]): P[A] = {
    table.foldLeft(term)(addPrecLevel)
  }

  def addPrecLevel[A](term: P[A], ops: List[Operator[A]]): P[A] = {
    val (ras, las, nas, prefix, postfix) = ops.foldRight(Batch.empty[A])(splitOp)
    val term1 = pTerm(choice(prefix), term, choice(postfix))
    val ras1 = pInfixR(choice(ras), term1)(_)
    val las1 = pInfixL(choice(las), term1)(_)
    val nas1 = pInfixN(choice(nas), term1)(_)

    term1 flatMap { x =>
      choice(List(ras1(x), las1(x), nas1(x), PassWith(x)))
    }
  }

  private object Batch {
    def empty[A]: Batch[A] = (Nil, Nil, Nil, Nil, Nil)
  }

  private type Batch[A] =
    ( List[P[A => A => A]]
    , List[P[A => A => A]]
    , List[P[A => A => A]]
    , List[P[A => A]]
    , List[P[A => A]]
    )

  private def splitOp[A](op: Operator[A], batch: Batch[A]): Batch[A] = {
    val (r, l, n, pre, post) = batch
    op match {
      case InfixR(op)  => (op :: r, l, n, pre, post)
      case InfixL(op)  => (r, op :: l, n, pre, post)
      case InfixN(op)  => (r, l, op :: n, pre, post)
      case Prefix(op)  => (r, l, n, op :: pre, post)
      case Postfix(op) => (r, l, n, pre, op :: post)
    }
  }

  private def pTerm[A](prefix: P[A => A], term: P[A], postfix: P[A => A]): P[A] = for {
    pre <- option(prefix)(x => x)
    x <- term
    post <- option(postfix)(x => x)
  } yield post(pre(x))

  private def pInfixN[A](op: P[A => A => A], p: P[A])(x: A): P[A] = for {
    f <- op
    y <- p
  } yield f(x)(y)

  private def pInfixL[A](op: P[A => A => A], p: P[A])(x: A): P[A] = for {
    f <- op
    y <- p
    r = f(x)(y)
    t <- pInfixL(op, p)(r) | PassWith(r)
  } yield t

  private def pInfixR[A](op: P[A => A => A], p: P[A])(x: A): P[A] = for {
    f <- op
    y <- p.flatMap(r => pInfixR(op, p)(r) | PassWith(r))
  } yield f(x)(y)

  private def option[A](p: P[A])(x: A): P[A] = {
    p | PassWith(x)
  }

  private def choice[A](ps: Seq[P[A]]): P[A] = {
    ps.foldRight[P[A]](Fail)(_ | _)
  }

}