2020VianExam

import Library._

object Desugar {

  def desugarProg(p: Prog): ExprC = p match {
    case ProgExt(ds, e) =>
      createDefs(ds, updateDefs(ds, desugar(e)))
  }

  // helper
  def getFunctionNames(ds: List[DefExt]): List[String] = ds match {
    case DefineExt(f, xs, expr) :: t => f :: getFunctionNames(t)
    case _ => Nil
  }

  def createDefs(ds: List[DefExt], e: ExprC): ExprC = {
    val fNames = getFunctionNames(ds)
    AppC(FdC(fNames, e), fNames.map(_ => UninitializedC()))
  }

  def updateDefs(ds: List[DefExt], e: ExprC): ExprC = ds match {
    case DefineExt(f, xs, expr) :: t => SeqC(SetC(f, FdC(xs, desugar(expr))), updateDefs(t, e))
    case _ => e

  }


  def desugar(e: ExprExt): ExprC = e match {
    case NumExt(n) => NumC(n)
    case PlusExt(e1, e2) => PlusC(desugar(e1), desugar(e2))
    case MinusExt(e1, e2) => MinusC(desugar(e1), desugar(e2))
    case MultExt(e1, e2) => MultC(desugar(e1), desugar(e2))
    case LtExt(e1, e2) => LtC(desugar(e1), desugar(e2))
    case IfExt(e, e1, e2) => IfC(desugar(e), desugar(e1), desugar(e2))
    case TrueExt() => TrueC()
    case FalseExt() => FalseC()
    case FdExt(xs, e) => FdC(xs, desugar(e))
    case AppExt(e, es) => AppC(desugar(e), es.map(desugar))
    case IdExt(x) => IdC(x)
    case SetExt(x, e) => SetC(x, desugar(e))
    case SeqExt(e1, e2) => SeqC(desugar(e1), desugar(e2))
  }

}


// DO NOT EDIT BELOW THIS LINE

sealed abstract class Prog
case class ProgExt(ds: List[DefExt], e: ExprExt)       extends Prog

sealed abstract class DefExt
case class DefineExt(f: String, xs: List[String], e: ExprExt) extends DefExt

sealed abstract class ExprExt
case class NumExt(n: Int)                              extends ExprExt
case class PlusExt(e1: ExprExt, e2: ExprExt)           extends ExprExt
case class MinusExt(e1: ExprExt, e2: ExprExt)          extends ExprExt
case class MultExt(e1: ExprExt, e2: ExprExt)           extends ExprExt
case class LtExt(e1: ExprExt, e2: ExprExt)             extends ExprExt
case class IfExt(e: ExprExt, e1: ExprExt, e2: ExprExt) extends ExprExt
case class TrueExt()                                   extends ExprExt
case class FalseExt()                                  extends ExprExt
case class FdExt(xs: List[String], e: ExprExt)         extends ExprExt
case class AppExt(e: ExprExt, es: List[ExprExt])       extends ExprExt
case class IdExt(x: String)                            extends ExprExt
case class SetExt(x: String, e: ExprExt)               extends ExprExt
case class SeqExt(e1: ExprExt, e2: ExprExt)            extends ExprExt

sealed abstract class ExprC
case class NumC(n: Int)                        extends ExprC
case class PlusC(e1: ExprC, e2: ExprC)         extends ExprC
case class MinusC(e1: ExprC, e2: ExprC)        extends ExprC
case class MultC(e1: ExprC, e2: ExprC)         extends ExprC
case class LtC(e1: ExprC, e2: ExprC)           extends ExprC
case class IfC(e: ExprC, e1: ExprC, e2: ExprC) extends ExprC
case class TrueC()                             extends ExprC
case class FalseC()                            extends ExprC
case class FdC(xs: List[String], e: ExprC)     extends ExprC
case class AppC(e: ExprC, es: List[ExprC])     extends ExprC
case class IdC(x: String)                      extends ExprC
case class SetC(x: String, e: ExprC)           extends ExprC
case class SeqC(e1: ExprC, e2: ExprC)          extends ExprC
case class UninitializedC()                    extends ExprC