cps_ir: Clean up BuilderTask.

tests/compiler/dart2js/js_backend_cps_ir_source_information_test.dart

Index: tests/compiler/dart2js/js_backend_cps_ir_source_information_test.dart
diff --git a/tests/compiler/dart2js/js_backend_cps_ir_source_information_test.dart b/tests/compiler/dart2js/js_backend_cps_ir_source_information_test.dart
index 505901f3d04d403be84d8dd7860be6f250ece2a0..713281f61146bfa9749a451faaca1098baef693e 100644
--- a/tests/compiler/dart2js/js_backend_cps_ir_source_information_test.dart
+++ b/tests/compiler/dart2js/js_backend_cps_ir_source_information_test.dart
@@ -43,10 +43,6 @@ js.Node getJsNodeForElement(Compiler compiler, Element element) {
   return compiler.enqueuer.codegen.generatedCode[element];
 }
 
-ir.FunctionDefinition getIrNodeForElement(Compiler compiler, Element element) {
-  return compiler.irBuilder.getIr(element);
-}
-
 String getCodeForMethod(Compiler compiler, String name) {
   Element foundElement;
   for (Element element in compiler.enqueuer.codegen.generatedCode.keys) {
@@ -71,13 +67,23 @@ runTests(List<TestEntry> tests) {
     Map files = {TEST_MAIN_FILE: test.source};
     asyncTest(() {
       Compiler compiler = compilerFor(files, options: <String>['--use-cps-ir']);
+      ir.FunctionDefinition irNodeForMain;
+
+      void cacheIrNodeForMain(Element function, ir.FunctionDefinition irNode) {
+        if (function == compiler.mainFunction) {
+          assert(irNodeForMain == null);
+          irNodeForMain = irNode;
+        }
+      }
+
       Uri uri = Uri.parse('memory:$TEST_MAIN_FILE');
+      compiler.irBuilder.builderCallback = cacheIrNodeForMain;
+
       return compiler.run(uri).then((bool success) {
         Expect.isTrue(success);
 
-        ir.Node irNode = getIrNodeForElement(compiler, compiler.mainFunction);
         IrSourceInformationVisitor irVisitor = new IrSourceInformationVisitor();
-        irNode.accept(irVisitor);
+        irNodeForMain.accept(irVisitor);
 
         js.Node jsNode = getJsNodeForElement(compiler, compiler.mainFunction);
         JsSourceInformationVisitor jsVisitor = new JsSourceInformationVisitor();
@@ -90,7 +96,7 @@ runTests(List<TestEntry> tests) {
               'Expected:\n$expectation\n'
               'but found\n${irVisitor.sourceInformation}\n'
               'in\n${test.source}'
-              'CPS:\n${irNode.accept(new ir.SExpressionStringifier())}');
+              'CPS:\n${irNodeForMain.accept(new ir.SExpressionStringifier())}');
         Expect.listEquals(expectation, jsVisitor.sourceInformation,
               'Unexpected JS source information. '
               'Expected:\n$expectation\n'