Free typechecker state and ASTs when they are no longer needed

mypy/build.py

@@ -59,6 +59,7 @@
 from mypy.typestate import TypeState, reset_global_state
 from mypy.renaming import VariableRenameVisitor
 from mypy.config_parser import parse_mypy_comments
+from mypy.freetree import free_tree
 
 
 # Switch to True to produce debug output related to fine-grained incremental
@@ -471,6 +472,7 @@ class BuildManager:
       missing_modules: Set of modules that could not be imported encountered so far
       stale_modules:   Set of modules that needed to be rechecked (only used by tests)
       fg_deps_meta:    Metadata for fine-grained dependencies caches associated with modules
+      fg_deps:         A fine-grained dependency map
       version_id:      The current mypy version (based on commit id when possible)
       plugin:          Active mypy plugin(s)
       plugins_snapshot:
@@ -523,6 +525,11 @@ def __init__(self, data_dir: str,
         self.modules = {}  # type: Dict[str, MypyFile]
         self.missing_modules = set()  # type: Set[str]
         self.fg_deps_meta = {}  # type: Dict[str, FgDepMeta]
+        # fg_deps holds the dependencies of every module that has been
+        # processed. We store this in BuildManager so that we can compute
+        # dependencies as we go, which allows us to free ASTs and type information,
+        # saving a ton of memory on net.
+        self.fg_deps = {}  # type: Dict[str, Set[str]]
         # Always convert the plugin to a ChainedPlugin so that it can be manipulated if needed
         if not isinstance(plugin, ChainedPlugin):
             plugin = ChainedPlugin(options, [plugin])
@@ -894,32 +901,24 @@ def invert_deps(deps: Dict[str, Set[str]],
     return rdeps
 
 
-def generate_deps_for_cache(proto_deps: Dict[str, Set[str]],
-                            manager: BuildManager,
+def generate_deps_for_cache(manager: BuildManager,
                             graph: Graph) -> Dict[str, Dict[str, Set[str]]]:
     """Generate fine-grained dependencies into a form suitable for serializing.
 
-    This does a few things:
-    1. Computes all fine grained deps from modules that were processed
-    2. Splits fine-grained deps based on the module of the trigger
-    3. For each module we generated fine-grained deps for, load any previous
+    This does a couple things:
+    1. Splits fine-grained deps based on the module of the trigger
+    2. For each module we generated fine-grained deps for, load any previous
        deps and merge them in.
 
     Returns a dictionary from module ids to all dependencies on that
     module. Dependencies not associated with a module in the build will be
     associated with the nearest parent module that is in the build, or the
     fake module FAKE_ROOT_MODULE if none are.
     """
-    from mypy.server.update import merge_dependencies  # Lazy import to speed up startup
-
-    # Compute the full set of dependencies from everything we've processed.
-    deps = {}  # type: Dict[str, Set[str]]
-    things = [st.compute_fine_grained_deps() for st in graph.values() if st.tree] + [proto_deps]
-    for st_deps in things:
-        merge_dependencies(st_deps, deps)
+    from mypy.server.deps import merge_dependencies  # Lazy import to speed up startup
 
     # Split the dependencies out into based on the module that is depended on.
-    rdeps = invert_deps(deps, graph)
+    rdeps = invert_deps(manager.fg_deps, graph)
 
     # We can't just clobber existing dependency information, so we
     # load the deps for every module we've generated new dependencies
@@ -2172,6 +2171,16 @@ def finish_passes(self) -> None:
                                 typemap=self.type_map())
             manager.report_file(self.tree, self.type_map(), self.options)
 
+            self.update_fine_grained_deps(self.manager.fg_deps)
+            self.free_state()
+            if not manager.options.fine_grained_incremental and not manager.options.preserve_asts:
+                free_tree(self.tree)
+
+    def free_state(self) -> None:
+        if self._type_checker:
+            self._type_checker.reset()
+            self._type_checker = None
+
     def _patch_indirect_dependencies(self,
                                      module_refs: Set[str],
                                      type_map: Dict[Expression, Type]) -> None:
@@ -2206,6 +2215,13 @@ def compute_fine_grained_deps(self) -> Dict[str, Set[str]]:
                                 python_version=self.options.python_version,
                                 options=self.manager.options)
 
+    def update_fine_grained_deps(self, deps: Dict[str, Set[str]]) -> None:
+        options = self.manager.options
+        if options.cache_fine_grained or options.fine_grained_incremental:
+            from mypy.server.deps import merge_dependencies  # Lazy import to speed up startup
+            merge_dependencies(self.compute_fine_grained_deps(), deps)
+            TypeState.update_protocol_deps(deps)
+
     def valid_references(self) -> Set[str]:
         assert self.ancestors is not None
         valid_refs = set(self.dependencies + self.suppressed + self.ancestors)
@@ -2616,10 +2632,9 @@ def dispatch(sources: List[BuildSource],
             # then we need to collect fine grained protocol dependencies.
             # Since these are a global property of the program, they are calculated after we
             # processed the whole graph.
-            TypeState.update_protocol_deps()
+            TypeState.add_all_protocol_deps(manager.fg_deps)
             if not manager.options.fine_grained_incremental:
-                proto_deps = TypeState.proto_deps or {}
-                rdeps = generate_deps_for_cache(proto_deps, manager, graph)
+                rdeps = generate_deps_for_cache(manager, graph)
                 write_deps_cache(rdeps, manager, graph)
 
     if manager.options.dump_deps:
@@ -3023,23 +3038,29 @@ def process_stale_scc(graph: Graph, scc: List[str], manager: BuildManager) -> No
             graph[id].semantic_analysis_pass_three()
         for id in stale:
             graph[id].semantic_analysis_apply_patches()
+
+    # Track what modules aren't yet done so we can finish them as soon
+    # as possible, saving memory.
+    unfinished_modules = set(stale)
     for id in stale:
         graph[id].type_check_first_pass()
-    more = True
-    while more:
-        more = False
+        if not graph[id].type_checker().deferred_nodes:
+            unfinished_modules.discard(id)
+            graph[id].finish_passes()
+
+    while unfinished_modules:
         for id in stale:
-            if graph[id].type_check_second_pass():
-                more = True
+            if id not in unfinished_modules:
+                continue
+            if not graph[id].type_check_second_pass():
+                unfinished_modules.discard(id)
+                graph[id].finish_passes()
     for id in stale:
         graph[id].generate_unused_ignore_notes()
     if any(manager.errors.is_errors_for_file(graph[id].xpath) for id in stale):
         for id in stale:
             graph[id].transitive_error = True
     for id in stale:
-        graph[id].finish_passes()
-        if manager.options.cache_fine_grained or manager.options.fine_grained_incremental:
-            graph[id].compute_fine_grained_deps()
         manager.flush_errors(manager.errors.file_messages(graph[id].xpath), False)
         graph[id].write_cache()
         graph[id].mark_as_rechecked()

mypy/freetree.py

@@ -0,0 +1,21 @@
+"""Generic node traverser visitor"""
+
+from mypy.traverser import TraverserVisitor
+from mypy.nodes import Block, MypyFile
+
+
+class TreeFreer(TraverserVisitor):
+    def visit_block(self, block: Block) -> None:
+        super().visit_block(block)
+        block.body.clear()
+
+
+def free_tree(tree: MypyFile) -> None:
+    """Free all the ASTs associated with a module.
+
+    This needs to be done recursively, since symbol tables contain
+    references to definitions, so those won't be freed but we want their
+    contents to be.
+    """
+    tree.accept(TreeFreer())
+    tree.defs.clear()

mypy/options.py

@@ -216,6 +216,12 @@ def __init__(self) -> None:
         # in modules being compiled. Not in the config file or command line.
         self.mypyc = False
 
+        # Disable the memory optimization of freeing ASTs when
+        # possible. This isn't exposed as a command line option
+        # because it is intended for software integrating with
+        # mypy. (Like mypyc.)
+        self.preserve_asts = False
+
         # Paths of user plugins
         self.plugins = []  # type: List[str]
 

mypy/server/deps.py

@@ -958,6 +958,12 @@ def visit_union_type(self, typ: UnionType) -> List[str]:
         return triggers
 
 
+def merge_dependencies(new_deps: Dict[str, Set[str]],
+                       deps: Dict[str, Set[str]]) -> None:
+    for trigger, targets in new_deps.items():
+        deps.setdefault(trigger, set()).update(targets)
+
+
 def non_trivial_bases(info: TypeInfo) -> List[TypeInfo]:
     return [base for base in info.mro[1:]
             if base.fullname() != 'builtins.object']

mypy/server/update.py

@@ -140,7 +140,7 @@
 from mypy.server.astmerge import merge_asts
 from mypy.server.aststrip import strip_target
 from mypy.server.aststripnew import strip_target_new, SavedAttributes
-from mypy.server.deps import get_dependencies_of_target
+from mypy.server.deps import get_dependencies_of_target, merge_dependencies
 from mypy.server.target import trigger_to_target
 from mypy.server.trigger import make_trigger, WILDCARD_TAG
 from mypy.util import module_prefix, split_target
@@ -163,7 +163,9 @@ def __init__(self, result: BuildResult) -> None:
         self.manager = manager
         self.graph = result.graph
         self.previous_modules = get_module_to_path_map(self.graph)
-        self.deps = get_all_dependencies(manager, self.graph)
+        self.deps = manager.fg_deps
+        # Merge in any root dependencies that may not have been loaded
+        merge_dependencies(manager.load_fine_grained_deps(FAKE_ROOT_MODULE), self.deps)
         self.previous_targets_with_errors = manager.errors.targets()
         self.previous_messages = result.errors[:]
         # Module, if any, that had blocking errors in the last run as (id, path) tuple.
@@ -378,7 +380,8 @@ def update_module(self,
             self.manager.log_fine_grained('triggered: %r' % sorted(filtered))
         self.triggered.extend(triggered | self.previous_targets_with_errors)
         if module in graph:
-            merge_dependencies(graph[module].compute_fine_grained_deps(), self.deps)
+            graph[module].update_fine_grained_deps(self.deps)
+            graph[module].free_state()
         remaining += propagate_changes_using_dependencies(
             manager, graph, self.deps, triggered,
             {module},
@@ -453,15 +456,11 @@ def ensure_trees_loaded(manager: BuildManager, graph: Dict[str, State],
         process_fresh_modules(graph, to_process, manager)
 
 
-def get_all_dependencies(manager: BuildManager, graph: Dict[str, State]) -> Dict[str, Set[str]]:
-    """Return the fine-grained dependency map for an entire build."""
-    # Deps for each module were computed during build() or loaded from the cache.
-    deps = manager.load_fine_grained_deps(FAKE_ROOT_MODULE)  # type: Dict[str, Set[str]]
-    for id in graph:
-        if graph[id].tree is not None:
-            merge_dependencies(graph[id].compute_fine_grained_deps(), deps)
-    TypeState.add_all_protocol_deps(deps)
-    return deps
+def fix_fg_dependencies(manager: BuildManager, deps: Dict[str, Set[str]]) -> None:
+    """Populate the dependencies with stuff that build may have missed"""
+    # This means the root module and typestate
+    merge_dependencies(manager.load_fine_grained_deps(FAKE_ROOT_MODULE), deps)
+    # TypeState.add_all_protocol_deps(deps)
 
 
 # The result of update_module_isolated when no blockers, with these items:
@@ -608,15 +607,12 @@ def restore(ids: List[str]) -> None:
     state.type_check_first_pass()
     state.type_check_second_pass()
     t2 = time.time()
-    state.compute_fine_grained_deps()
-    t3 = time.time()
     state.finish_passes()
-    t4 = time.time()
+    t3 = time.time()
     manager.add_stats(
         semanal_time=t1 - t0,
         typecheck_time=t2 - t1,
-        deps_time=t3 - t2,
-        finish_passes_time=t4 - t3)
+        finish_passes_time=t3 - t2)
 
     graph[module] = state
 
@@ -700,14 +696,6 @@ def get_sources(fscache: FileSystemCache,
     return sources
 
 
-def merge_dependencies(new_deps: Dict[str, Set[str]],
-                       deps: Dict[str, Set[str]]) -> None:
-    for trigger, targets in new_deps.items():
-        deps.setdefault(trigger, set()).update(targets)
-    # Merge also the newly added protocol deps.
-    TypeState.update_protocol_deps(deps)
-
-
 def calculate_active_triggers(manager: BuildManager,
                               old_snapshots: Dict[str, Dict[str, SnapshotItem]],
                               new_modules: Dict[str, Optional[MypyFile]]) -> Set[str]:
@@ -994,6 +982,8 @@ def key(node: FineGrainedDeferredNode) -> int:
     # Report missing imports.
     graph[module_id].verify_dependencies()
 
+    graph[module_id].free_state()
+
     return new_triggered
 
 

mypy/test/testdeps.py

@@ -45,6 +45,7 @@ def run_case(self, testcase: DataDrivenTestCase) -> None:
         options.cache_dir = os.devnull
         options.python_version = python_version
         options.export_types = True
+        options.preserve_asts = True
         messages, files, type_map = self.build(src, options)
         a = messages
         if files is None or type_map is None:

mypy/test/testmerge.py

@@ -106,6 +106,7 @@ def build(self, source: str, testcase: DataDrivenTestCase) -> Optional[BuildResu
         options.incremental = True
         options.fine_grained_incremental = True
         options.use_builtins_fixtures = True
+        options.export_types = True
         options.show_traceback = True
         options.python_version = PYTHON3_VERSION
         main_path = os.path.join(test_temp_dir, 'main')
@@ -216,7 +217,13 @@ def dump_types(self, manager: FineGrainedBuildManager) -> List[str]:
         for module_id in sorted(manager.manager.modules):
             if not is_dumped_module(module_id):
                 continue
-            type_map = manager.graph[module_id].type_map()
+            all_types = manager.manager.all_types
+            # Compute a module type map from the global type map
+            tree = manager.graph[module_id].tree
+            assert tree is not None
+            type_map = {node: all_types[node]
+                        for node in get_subexpressions(tree)
+                        if node in all_types}
             if type_map:
                 a.append('## {}'.format(module_id))
                 for expr in sorted(type_map, key=lambda n: (n.line, short_type(n),

mypy/test/testtypegen.py

@@ -31,6 +31,7 @@ def run_case(self, testcase: DataDrivenTestCase) -> None:
             options.use_builtins_fixtures = True
             options.show_traceback = True
             options.export_types = True
+            options.preserve_asts = True
             result = build.build(sources=[BuildSource('main', None, src)],
                                  options=options,
                                  alt_lib_path=test_temp_dir)

mypy/typestate.py

@@ -211,15 +211,13 @@ def update_protocol_deps(second_map: Optional[Dict[str, Set[str]]] = None) -> No
     def add_all_protocol_deps(deps: Dict[str, Set[str]]) -> None:
         """Add all known protocol dependencies to deps.
 
-        This is used by tests and debug output, and also when passing
-        all collected or loaded dependencies on to FineGrainedBuildManager
-        in its __init__.
+        This is used by tests and debug output, and also when collecting
+        all collected or loaded dependencies as part of build.
         """
         TypeState.update_protocol_deps()  # just in case
-        assert TypeState.proto_deps is not None, (
-            "This should not be called after failed cache load")
-        for trigger, targets in TypeState.proto_deps.items():
-            deps.setdefault(trigger, set()).update(targets)
+        if TypeState.proto_deps is not None:
+            for trigger, targets in TypeState.proto_deps.items():
+                deps.setdefault(trigger, set()).update(targets)
 
 
 def reset_global_state() -> None: