diff --git a/packages/devtools_app/lib/src/screens/app_size/app_size_controller.dart b/packages/devtools_app/lib/src/screens/app_size/app_size_controller.dart
index c8f8669..c7926f2 100644
--- a/packages/devtools_app/lib/src/screens/app_size/app_size_controller.dart
+++ b/packages/devtools_app/lib/src/screens/app_size/app_size_controller.dart
@@ -17,7 +17,7 @@
 import 'app_size_screen.dart';
 
 // Temporary feature flag for deferred loading.
-bool deferredLoadingSupportEnabled = false;
+bool deferredLoadingSupportEnabled = true;
 
 const _artificialRootNodeName = 'ArtificialRoot';
 const _entireAppNodeName = 'Entire App';
@@ -61,6 +61,20 @@
   }
 }
 
+class DiffTreeMap {
+  const DiffTreeMap({
+    required this.combined,
+    required this.increaseOnly,
+    required this.decreaseOnly,
+  });
+
+  final TreemapNode? combined;
+
+  final TreemapNode? increaseOnly;
+
+  final TreemapNode? decreaseOnly;
+}
+
 class AppSizeController {
   static const unsupportedFileTypeError =
       'Failed to load size analysis file: file type not supported.\n\n'
@@ -252,9 +266,13 @@
   ValueListenable<AppUnit> get selectedAppUnit => _selectedAppUnit;
   final _selectedAppUnit = ValueNotifier<AppUnit>(AppUnit.entireApp);
 
-  void changeSelectedAppUnit(AppUnit appUnit) {
+  void changeSelectedAppUnit(AppUnit appUnit, Key tabKey) {
     _selectedAppUnit.value = appUnit;
-    _loadApp(_dataForAppUnit!);
+    if (tabKey == AppSizeScreen.analysisTabKey) {
+      _loadApp(_dataForAppUnit!);
+    } else {
+      print('selected $appUnit in diff view');
+    }
   }
 
   /// Notifies that the json files are currently being processed.
@@ -343,7 +361,8 @@
   ) {
     if (_hasDeferredInfo(jsonFile)) {
       jsonFile['children'] = _extractChildren(jsonFile)
-          .where((child) => child['isDeferred'] == true);
+          .where((child) => child['isDeferred'] == true)
+          .toList();
       jsonFile['n'] = _deferredNodeName;
     }
     return jsonFile;
@@ -385,55 +404,52 @@
     await delayForBatchProcessing(micros: 10000);
 
     Map<String, dynamic> diffMap;
-    if (oldFile.isAnalyzeSizeFile && newFile.isAnalyzeSizeFile) {
-      var oldFileJson = oldFile.data as Map<String, dynamic>;
-      var newFileJson = newFile.data as Map<String, dynamic>;
+    Map<String, dynamic> mainDiffMap;
+    Map<String, dynamic> deferredDiffMap;
 
-      if (_hasDeferredInfo(oldFileJson) || _hasDeferredInfo(newFileJson)) {
+    if (oldFile.isAnalyzeSizeFile && newFile.isAnalyzeSizeFile) {
+      final oldFileJson = oldFile.data as Map<String, dynamic>;
+      final newFileJson = newFile.data as Map<String, dynamic>;
+
+      if (!_hasDeferredInfo(oldFileJson) && !_hasDeferredInfo(newFileJson)) {
+        diffMap = _generateDiffMapFromAnalyzeSizeFiles(
+          oldFileJson: oldFileJson,
+          newFileJson: newFileJson,
+        );
+      } else {
         _isDeferredApp.value = deferredLoadingSupportEnabled;
+        Map<String, dynamic> oldEntireAppFileJson = oldFileJson;
+        Map<String, dynamic> newEntireAppFileJson = newFileJson;
 
         if (!_hasDeferredInfo(oldFileJson)) {
-          oldFileJson = _wrapInArtificialRoot(oldFileJson);
+          oldEntireAppFileJson = _wrapInArtificialRoot(oldFileJson);
         } else if (!_hasDeferredInfo(newFileJson)) {
-          newFileJson = _wrapInArtificialRoot(newFileJson);
+          newEntireAppFileJson = _wrapInArtificialRoot(newFileJson);
         }
-      }
 
-      final oldApkProgramInfo = ProgramInfo();
-      _apkJsonToProgramInfo(
-        program: oldApkProgramInfo,
-        parent: oldApkProgramInfo.root,
-        json: oldFileJson,
-      );
+        final oldMainOnlyFileJson = _extractMainUnit(oldEntireAppFileJson);
+        final newMainOnlyFileJson = _extractMainUnit(newEntireAppFileJson);
 
-      // Extract the precompiler trace from the old file, if it exists, and
-      // generate a call graph.
-      final oldPrecompilerTrace = oldFileJson.remove('precompiler-trace');
-      if (oldPrecompilerTrace != null) {
-        _oldDiffCallGraph = generateCallGraphWithDominators(
-          oldPrecompilerTrace,
-          NodeType.packageNode,
+        final oldDeferredOnlyFileJson =
+            _extractDeferredUnits(oldEntireAppFileJson);
+        final newDeferredOnlyFileJson =
+            _extractDeferredUnits(newEntireAppFileJson);
+
+        diffMap = _generateDiffMapFromAnalyzeSizeFiles(
+          oldFileJson: oldEntireAppFileJson,
+          newFileJson: newEntireAppFileJson,
+        );
+
+        mainDiffMap = _generateDiffMapFromAnalyzeSizeFiles(
+          oldFileJson: oldMainOnlyFileJson,
+          newFileJson: newMainOnlyFileJson,
+        );
+
+        deferredDiffMap = _generateDiffMapFromAnalyzeSizeFiles(
+          oldFileJson: oldDeferredOnlyFileJson,
+          newFileJson: newDeferredOnlyFileJson,
         );
       }
-
-      final newApkProgramInfo = ProgramInfo();
-      _apkJsonToProgramInfo(
-        program: newApkProgramInfo,
-        parent: newApkProgramInfo.root,
-        json: newFileJson,
-      );
-
-      // Extract the precompiler trace from the new file, if it exists, and
-      // generate a call graph.
-      final newPrecompilerTrace = newFileJson.remove('precompiler-trace');
-      if (newPrecompilerTrace != null) {
-        _newDiffCallGraph = generateCallGraphWithDominators(
-          newPrecompilerTrace,
-          NodeType.packageNode,
-        );
-      }
-
-      diffMap = compareProgramInfo(oldApkProgramInfo, newApkProgramInfo);
     } else {
       try {
         diffMap = buildComparisonTreemap(oldFile.data, newFile.data);
@@ -458,27 +474,78 @@
     diffMap['n'] = isDeferredApp.value ? _entireAppNodeName : _rootNodeName;
 
     // TODO(peterdjlee): Try to move the non-active tree generation to separate isolates.
-    _combinedDiffTreeRoot = generateDiffTree(
-      diffMap,
-      DiffTreeType.combined,
-      skipNodesWithNoByteSizeChange: !isDeferredApp.value,
-    );
-    _increasedDiffTreeRoot = generateDiffTree(
-      diffMap,
-      DiffTreeType.increaseOnly,
-      skipNodesWithNoByteSizeChange: !isDeferredApp.value,
-    );
-    _decreasedDiffTreeRoot = generateDiffTree(
-      diffMap,
-      DiffTreeType.decreaseOnly,
-      skipNodesWithNoByteSizeChange: !isDeferredApp.value,
-    );
+    final diffTreeMap = _generateDiffTrees(diffMap);
+    _combinedDiffTreeRoot = diffTreeMap.combined;
+    _increasedDiffTreeRoot = diffTreeMap.increaseOnly;
+    _decreasedDiffTreeRoot = diffTreeMap.decreaseOnly;
 
     changeDiffRoot(_activeDiffRoot);
 
     _processingNotifier.value = false;
   }
 
+  DiffTreeMap _generateDiffTrees(Map<String, dynamic> diffMap) {
+    // TODO(peterdjlee): Try to move the non-active tree generation to separate isolates.
+    return DiffTreeMap(
+      combined: generateDiffTree(
+        diffMap,
+        DiffTreeType.combined,
+        skipNodesWithNoByteSizeChange: !isDeferredApp.value,
+      ),
+      increaseOnly: generateDiffTree(
+        diffMap,
+        DiffTreeType.increaseOnly,
+        skipNodesWithNoByteSizeChange: !isDeferredApp.value,
+      ),
+      decreaseOnly: generateDiffTree(
+        diffMap,
+        DiffTreeType.decreaseOnly,
+        skipNodesWithNoByteSizeChange: !isDeferredApp.value,
+      ),
+    );
+  }
+
+  Map<String, dynamic> _generateDiffMapFromAnalyzeSizeFiles({
+    required Map<String, dynamic> oldFileJson,
+    required Map<String, dynamic> newFileJson,
+  }) {
+    final oldApkProgramInfo = ProgramInfo();
+    _apkJsonToProgramInfo(
+      program: oldApkProgramInfo,
+      parent: oldApkProgramInfo.root,
+      json: oldFileJson,
+    );
+
+    // Extract the precompiler trace from the old file, if it exists, and
+    // generate a call graph.
+    final oldPrecompilerTrace = oldFileJson.remove('precompiler-trace');
+    if (oldPrecompilerTrace != null) {
+      _oldDiffCallGraph = generateCallGraphWithDominators(
+        oldPrecompilerTrace,
+        NodeType.packageNode,
+      );
+    }
+
+    final newApkProgramInfo = ProgramInfo();
+    _apkJsonToProgramInfo(
+      program: newApkProgramInfo,
+      parent: newApkProgramInfo.root,
+      json: newFileJson,
+    );
+
+    // Extract the precompiler trace from the new file, if it exists, and
+    // generate a call graph.
+    final newPrecompilerTrace = newFileJson.remove('precompiler-trace');
+    if (newPrecompilerTrace != null) {
+      _newDiffCallGraph = generateCallGraphWithDominators(
+        newPrecompilerTrace,
+        NodeType.packageNode,
+      );
+    }
+
+    return compareProgramInfo(oldApkProgramInfo, newApkProgramInfo);
+  }
+
   Map<String, dynamic> _wrapInArtificialRoot(Map<String, dynamic> json) {
     json['n'] = _mainNodeName;
     return <String, dynamic>{
diff --git a/packages/devtools_app/lib/src/screens/app_size/app_size_screen.dart b/packages/devtools_app/lib/src/screens/app_size/app_size_screen.dart
index 8ad5fb9..211c38a 100644
--- a/packages/devtools_app/lib/src/screens/app_size/app_size_screen.dart
+++ b/packages/devtools_app/lib/src/screens/app_size/app_size_screen.dart
@@ -208,9 +208,7 @@
                   ),
                   Row(
                     children: [
-                      if (currentTab.key == AppSizeScreen.analysisTabKey &&
-                          isDeferredApp)
-                        _buildAppUnitDropdown(),
+                      if (isDeferredApp) _buildAppUnitDropdown(currentTab.key!),
                       if (currentTab.key == AppSizeScreen.diffTabKey)
                         _buildDiffTreeTypeDropdown(),
                       const SizedBox(width: defaultSpacing),
@@ -253,7 +251,7 @@
     );
   }
 
-  DropdownButtonHideUnderline _buildAppUnitDropdown() {
+  DropdownButtonHideUnderline _buildAppUnitDropdown(Key tabKey) {
     return DropdownButtonHideUnderline(
       key: AppSizeScreen.appUnitDropdownKey,
       child: DropdownButton<AppUnit>(
@@ -265,7 +263,7 @@
         ],
         onChanged: (newAppUnit) {
           setState(() {
-            controller.changeSelectedAppUnit(newAppUnit!);
+            controller.changeSelectedAppUnit(newAppUnit!, tabKey);
           });
         },
       ),