Store ResolvedAst on AstElement

pkg/compiler/lib/src/serialization/resolved_ast_serialization.dart

Index: pkg/compiler/lib/src/serialization/resolved_ast_serialization.dart
diff --git a/pkg/compiler/lib/src/serialization/resolved_ast_serialization.dart b/pkg/compiler/lib/src/serialization/resolved_ast_serialization.dart
index de45dd1314fbcf58968e8e23346f6bb0bd3b70f2..9b4e0d363455a39055f85af7a9e6dcc905dc0e51 100644
--- a/pkg/compiler/lib/src/serialization/resolved_ast_serialization.dart
+++ b/pkg/compiler/lib/src/serialization/resolved_ast_serialization.dart
@@ -25,6 +25,7 @@ import 'keys.dart';
 import 'modelz.dart';
 import 'serialization.dart';
 import 'serialization_util.dart';
+import 'modelz.dart';
 
 /// Visitor that computes a node-index mapping.
 class AstIndexComputer extends Visitor {
@@ -344,18 +345,18 @@ class ResolvedAstDeserializer {
       ObjectDecoder objectDecoder,
       ParsingContext parsing,
       Token getBeginToken(Uri uri, int charOffset),
-      DeserializerPlugin nativeDataDeserializer,
-      Map<Element, ResolvedAst> resolvedAstMap) {
+      DeserializerPlugin nativeDataDeserializer) {
     ResolvedAstKind kind =
         objectDecoder.getEnum(Key.KIND, ResolvedAstKind.values);
     switch (kind) {
       case ResolvedAstKind.PARSED:
         deserializeParsed(element, objectDecoder, parsing, getBeginToken,
-            nativeDataDeserializer, resolvedAstMap);
+            nativeDataDeserializer);
         break;
       case ResolvedAstKind.DEFAULT_CONSTRUCTOR:
       case ResolvedAstKind.FORWARDING_CONSTRUCTOR:
-        resolvedAstMap[element] = new SynthesizedResolvedAst(element, kind);
+        (element as AstElementMixinZ).resolvedAst =
+            new SynthesizedResolvedAst(element, kind);
         break;
     }
   }
@@ -364,12 +365,11 @@ class ResolvedAstDeserializer {
   /// method, or field) and its nested closures. The [ResolvedAst]s are added
   /// to [resolvedAstMap].
   static void deserializeParsed(
-      Element element,
+      AstElementMixinZ element,
       ObjectDecoder objectDecoder,
       ParsingContext parsing,
       Token getBeginToken(Uri uri, int charOffset),
-      DeserializerPlugin nativeDataDeserializer,
-      Map<Element, ResolvedAst> resolvedAstMap) {
+      DeserializerPlugin nativeDataDeserializer) {
     CompilationUnitElement compilationUnit = element.compilationUnit;
     DiagnosticReporter reporter = parsing.reporter;
     Uri uri = objectDecoder.getUri(Key.URI);
@@ -480,7 +480,7 @@ class ResolvedAstDeserializer {
               builder.emptyStatement());
           return constructorNode;
         case AstKind.ENUM_CONSTANT:
-          EnumConstantElement enumConstant = element;
+          EnumConstantElementZ enumConstant = element;
           EnumClassElement enumClass = element.enclosingClass;
           int index = enumConstant.index;
           AstBuilder builder = new AstBuilder(element.sourcePosition.begin);
@@ -660,13 +660,11 @@ class ResolvedAstDeserializer {
             elements.registerNativeData(node, nativeData);
           }
         }
-        FunctionElement function =
+        LocalFunctionElementZ function =
             objectDecoder.getElement(Key.FUNCTION, isOptional: true);
         if (function != null) {
           FunctionExpression functionExpression = node;
-          assert(invariant(function, !resolvedAstMap.containsKey(function),
-              message: "ResolvedAst has already been computed for $function."));
-          resolvedAstMap[function] = new ParsedResolvedAst(function,
+          function.resolvedAst = new ParsedResolvedAst(function,
               functionExpression, functionExpression.body, elements, uri);
         }
         // TODO(johnniwinther): Remove these when inference doesn't need `.node`
@@ -683,9 +681,7 @@ class ResolvedAstDeserializer {
         }
       }
     }
-    assert(invariant(element, !resolvedAstMap.containsKey(element),
-        message: "ResolvedAst has already been computed for $element."));
-    resolvedAstMap[element] =
+    element.resolvedAst =
         new ParsedResolvedAst(element, root, body, elements, uri);
   }
 }