// Source language: arithmetic expressions with function calls 

enum Expr {

    case Num(n: Int)

    case Aop(op: String, e1: Expr, e2: Expr)

    case Call(fname: String, args: List[Expr])

}

import Expr._

// Target language 

// "trivial" KValues

enum KVal {

    case KVar(s: String)

    case KNum(n: Int)

    case KAop(op: String, v1: KVal, v2: KVal)

    case KCall(fname: String, args: List[KVal])

}

import KVal._

// KExpressions 

enum KExp {

    case KReturn(v: KVal)

    case KLet(x: String, v: KVal, e: KExp) 

}

import KExp._

// Intermediate KExpressions2 

enum KExp2 {

    case KPrim(v: KVal)

    case KLet2(x: String, e1: KExp2, e2: KExp2) 

}

import KExp2._

def pexp(e: KExp): String = e match {

    case KReturn(v) => s"KReturn($v)"

    case KLet(x,e1,e2) => s"KLet($x = ${e1} \n in ${pexp(e2)})"

}

var cnt = -1

def Fresh(s: String) = {

    cnt = cnt + 1

    s"${s}_${cnt}"

}

def KNorm(e: Expr): KExp2 = e match { 

    case Num(i) => KPrim(KNum(i))

    case Aop(op, l, r) => {

        val zl = Fresh("zl")

        val zr = Fresh("zr")

        KLet2(zl, KNorm(l), 

           KLet2(zr, KNorm(r), 

              KPrim(KAop(op, KVar(zl), KVar(zr)))))

    }

    case Call(fname, args) => {

        def aux(args: List[Expr], zs: List[KVal]) : KExp2 = args match {

            case Nil => KPrim(KCall(fname, zs))

            case a::as => {

                val z = Fresh("z")

                KLet2(z, KNorm(a), aux(as, zs ::: List(KVar(z))))

            }

        }

        aux(args, Nil)  

    }

}

def KNormi(e: Expr) = {

   val z = Fresh("z")

   KLet2(z, KNorm(e), KPrim(KVar(z)))

}

def Kdenest(ke: KExp2) : KExp = ke match {

    case KPrim(v) => KReturn(v)

    case KLet2(xt, e1, e2) => {

        def insert(ke: KExp) : KExp = ke match {

            case KLet(yt, e3, e4) => KLet(yt, e3, insert(e4))

            case KReturn(v) => KLet(xt, v, Kdenest(e2))

        }

        insert(Kdenest(e1))

    }

}

// Constant Folding

def opt(v: KVal, env: Map[String, Int]) : KVal = v match {

    case KVar(s) => if (env.isDefinedAt(s)) KNum(env(s)) else KVar(s)

    case KNum(n) => KNum(n)

    case KAop(op, v1, v2) => (op, opt(v1, env), opt(v2, env)) match {

        case ("+", KNum(n1), KNum(n2)) => KNum(n1 + n2)

        case ("*", KNum(n1), KNum(n2)) => KNum(n1 * n2)

        case (_, v1o, v2o) => KAop(op, v1o, v2o)

    }

    case KCall(fname, args) => KCall(fname, args.map(opt(_, env)))

}

def Koptimise(ke: KExp, env: Map[String, Int] = Map()) : KExp = ke match {

    case KReturn(v) => KReturn(opt(v, env)) 

    case KLet(x, v, e) => opt(v, env) match {

        case KNum(n) => Koptimise(e, env + (x -> n))

        case vo => KLet(x, vo, Koptimise(e, env))

    }

}    

//1 + foo(bar(4 * -7), 3, id(12))

val etest = 

    Aop("+", Num(1),

             Call("foo", 

                List(Call("bar", 

                             List(Aop("*", Num(4), Num(-7)))), 

                     Num(3), 

                     Call("id", List(Num(12))))))

val ke = Koptimise(Kdenest(KNormi(etest)))

println(pexp(ke))