fix cpshighoptimizer (tests pass)

compiler/src/l3/CPSOptimizer.scala

@@ -55,74 +55,71 @@ abstract class CPSOptimizer[T <: CPSTreeModule { type Name = Symbol }]
   private def shrink(tree: Tree): Tree =
     shrink(tree, State(census(tree)))
 
-  private def replaceArgs(args: Seq[Atom], s: State): Seq[Atom] = args map { a =>
-    s.aSubst.getOrElse(a, a)
-  }
+  private def replaceArgs(args: Seq[Atom], s: State): Seq[Atom] = 
+    args map { a => s.aSubst.getOrElse(a, a) }
+  
   private def shrink(tree: Tree, s: State): Tree = tree match {
     
-    case AppC(cntName, args) =>
+    case AppC(oldCntName, args) => {
       val replacedArgs = replaceArgs(args, s)
-      if (s.appliedOnce(cntName)) {
-        // Inline
+      val cntName = s.cSubst.getOrElse(oldCntName, oldCntName)
+
+      if (s.appliedOnce(cntName) && s.cEnv.contains(cntName)) {
         val cnt = s.cEnv(cntName)
         val newState = s.withASubst(cnt.args, replacedArgs)
         shrink(cnt.body, newState)
       } else {
         AppC(cntName, replacedArgs)
       }
-    case AppF(funAtom, retC, args) =>
+    }
+      
+    case AppF(oldFunAtom, oldRetC, args) => {
+      val funAtom = s.aSubst.getOrElse(oldFunAtom, oldFunAtom)
       val replacedArgs = replaceArgs(args, s)
+      val retC = s.cSubst.getOrElse(oldRetC, oldRetC)
+
       funAtom match {
         case AtomN(n) if s.fEnv.contains(n) && s.appliedOnce(n) => 
-          // Inline
           val fun = s.fEnv(n)
-          val newState = s.withASubst(fun.args, replacedArgs)
-                          .withCSubst(fun.retC, retC)
-          val newBody = shrink(fun.body, newState)
-          newBody
+          val newState = s.withASubst(fun.args, replacedArgs).withCSubst(fun.retC, retC)
+          shrink(fun.body, newState)
         case _ => AppF(funAtom, retC, replacedArgs)
       }
+    }
 
-    case LetF(funs, body) =>
-      val undeadFuns = funs.filter(fun => !s.dead(fun.name))
-      val undeadShrunkFuns = undeadFuns.map { fun =>
-        Fun(fun.name, fun.retC, fun.args, shrink(fun.body, s))}
-      val nonInlined = undeadShrunkFuns
-        .filter(fun => !s.appliedOnce(fun.name))
-
-      val newState = s.withFuns(undeadShrunkFuns)
-      val newBody = shrink(body, newState)
-      if (nonInlined.isEmpty) {
-        newBody
-      } else {
-        LetF(nonInlined, newBody)
+    case LetF(funs, body) =>  {
+      val undeadFuns = funs.filterNot(f => s.dead(f.name))
+      val nonInlined = undeadFuns.filter(fun => !s.appliedOnce(fun.name))
+      val newState = s.withFuns(undeadFuns)
+
+      val shrunkFuns = nonInlined.map(f => Fun(f.name, f.retC, f.args, shrink(f.body, newState)))
+      val shrunkBody = shrink(body, newState)
+
+      if (shrunkFuns.isEmpty) shrunkBody 
+      else LetF(shrunkFuns, shrunkBody)
     }
     
-    case LetP(name, prim, args, body) =>
-      if (s.dead(name) && !impure(prim)) {
-        shrink(body, s)
-      } else {
+    case LetP(name, prim, args, body) => {
       val replacedArgs = replaceArgs(args, s)
-        val primArgPair = (prim, replacedArgs)
-        if (!unstable(prim) && !impure(prim) && s.eInvEnv.contains(primArgPair)) {
-          val newS = s.withASubst(name, s.eInvEnv(primArgPair))
-          val newBody = shrink(body, newS)
-          newBody
+      val newState = s.withExp(name, prim, replacedArgs)
+      val shrunkBody = shrink(body, newState)
+
+      if (s.dead(name) && !impure(prim)) {
+        shrunkBody
       } else {
-          val newS = s.withExp(name, prim, replacedArgs)
-          LetP(name, prim, replacedArgs, shrink(body, newS))
+        LetP(name, prim, replacedArgs, shrunkBody)
       }
     }
 
-    case LetC(cnts, body) =>
-      val undeadConts = cnts.filter(cnt => !s.dead(cnt.name))
+    case LetC(cnts, body) => {
+      val undeadConts = cnts.filterNot(cnt => s.dead(cnt.name))
       val undeadShrunkConts = undeadConts.map { cnt =>
         val newBody = shrink(cnt.body, s)
         Cnt(cnt.name, cnt.args, newBody)
       }
 
-      val nonInlinedConts = undeadShrunkConts.filter { cnt =>
-        !s.appliedOnce(cnt.name)
+      val nonInlinedConts = undeadShrunkConts.filterNot { cnt =>
+        s.appliedOnce(cnt.name)
       }
 
       val newBody = shrink(body, s.withCnts(undeadShrunkConts))
@@ -131,6 +128,10 @@ abstract class CPSOptimizer[T <: CPSTreeModule { type Name = Symbol }]
       } else {
         LetC(nonInlinedConts, newBody)
       }
+    }
+
+    case If(cond, args, thenC, elseC) => 
+      If(cond, replaceArgs(args, s), thenC, elseC)
 
     case Halt(a) => 
       Halt(s.aSubst.getOrElse(a, a))