Introduce ModuleAnswersContext

Summary:
# Context
This stack will re-design how `pyrefly --report-pysa` is implemented. Right now, it is a post processing step that requires pyrefly to preserve internal states for all modules, which can lead to high memory usage.
To avoid this, we will incorporate the pysa reporting steps within the type checking.
See https://docs.google.com/document/d/1Bk8izFv4nQxbQmaog9OqcpcSxyzAW3Z-bn0JUV4wkg4/ for context

# This diff
In the next diffs, we will start building pysa information while type checking is running. To do this properly, we need to be integrated within the different resolution steps (ast, bindings, answers, solutions, etc.) of pyrefly. The obvious approach so to build pysa information once solutions for a module are built, since that's when we have all the information about a module.

To do this properly, we should ensure that our pysa information building step avoids querying for cross-module information, since that creates cycles during the resolution.

One solution is to avoid accessing the `transaction` when building this information. To do this, we introduce a `ModuleAnswersContext` struct from `ModuleContext`, containing the
per-module data that doesn't require cross-module access (handle, module_id,
module_info, stdlib, ast, bindings, answers). `ModuleContext` now wraps
ModuleAnswersContext alongside `module_ids` and `transaction` which allow cross-module information access.

This diff also inlines `get_class_field_from_current_class_only` to avoid `ad_hoc_solve`
for current-module class field lookups, narrowing it to take `ModuleAnswersContext` instead of `ModuleContext`.

Reviewed By: tianhan0

Differential Revision: D97328329

fbshipit-source-id: 02e70f6967f21a6d4cc92a671adf434759e63ced

Author: arthaud

pyrefly/lib/report/pysa/ast_visitor.rs

@@ -307,20 +307,24 @@ fn visit_statement<V: AstScopedVisitor>(
         Stmt::FunctionDef(function_def) => {
             let key = KeyDecoratedFunction(ShortIdentifier::new(&function_def.name));
             let function_scope = if let Some(idx) = module_context
+                .answers_context
                 .bindings
                 .key_to_idx_hashed_opt(Hashed::new(&key))
             {
                 let decorated_function = DecoratedFunction::from_bindings_answers(
                     idx,
-                    &module_context.bindings,
-                    &module_context.answers,
+                    &module_context.answers_context.bindings,
+                    &module_context.answers_context.answers,
                 );
-                if should_export_decorated_function(&decorated_function, module_context) {
+                if should_export_decorated_function(
+                    &decorated_function,
+                    &module_context.answers_context,
+                ) {
                     Scope::ExportedFunction {
                         function_id: FunctionId::Function {
                             location: PysaLocation::from_text_range(
                                 function_def.identifier().range(),
-                                &module_context.module_info,
+                                &module_context.answers_context.module_info,
                             ),
                         },
                         location: function_def.identifier().range(),
@@ -424,10 +428,12 @@ fn visit_statement<V: AstScopedVisitor>(
         Stmt::ClassDef(class_def) => {
             let key = KeyClass(ShortIdentifier::new(&class_def.name));
             let class_scope = if let Some(idx) = module_context
+                .answers_context
                 .bindings
                 .key_to_idx_hashed_opt(Hashed::new(&key))
             {
                 let class = module_context
+                    .answers_context
                     .answers
                     .get_idx(idx)
                     .unwrap()
@@ -532,11 +538,11 @@ pub fn visit_module_ast<V: AstScopedVisitor>(
     let mut scopes = Scopes {
         stack: vec![Scope::TopLevel],
     };
-    visitor.enter_toplevel_scope(&module_context.ast, &scopes);
+    visitor.enter_toplevel_scope(&module_context.answers_context.ast, &scopes);
     visitor.on_scope_update(&scopes);
-    for stmt in &module_context.ast.body {
+    for stmt in &module_context.answers_context.ast.body {
         visit_statement(stmt, visitor, &mut scopes, module_context);
     }
-    visitor.exit_toplevel_scope(&module_context.ast, &scopes);
+    visitor.exit_toplevel_scope(&module_context.answers_context.ast, &scopes);
     scopes
 }

pyrefly/lib/report/pysa/call_graph.rs

@@ -121,11 +121,12 @@ impl<'a> DefinitionToFunctionMapVisitor<'a> {
     fn bind_name(&mut self, key: Key, scopes: &Scopes) {
         if let Some(idx) = self
             .module_context
+            .answers_context
             .bindings
             .key_to_idx_hashed_opt(Hashed::new(&key))
             && let Some(current_function) = scopes.current_exported_function(
-                self.module_context.module_id,
-                self.module_context.module_info.name(),
+                self.module_context.answers_context.module_id,
+                self.module_context.answers_context.module_info.name(),
                 &SCOPE_EXPORTED_FUNCTION_FLAGS,
             )
         {
@@ -220,9 +221,10 @@ impl<'a> CapturedVariableVisitor<'a> {
     fn get_definition_from_usage(&self, key: Key) -> Option<FunctionRef> {
         let idx = self
             .module_context
+            .answers_context
             .bindings
             .key_to_idx_hashed_opt(Hashed::new(&key))?;
-        let binding = self.module_context.bindings.get(idx);
+        let binding = self.module_context.answers_context.bindings.get(idx);
         match binding {
             Binding::Forward(definition_idx) | Binding::ForwardToFirstUse(definition_idx) => {
                 self.get_definition_from_idx(
@@ -257,7 +259,7 @@ impl<'a> CapturedVariableVisitor<'a> {
         }
         depth += 1;
 
-        let binding = self.module_context.bindings.get(idx);
+        let binding = self.module_context.answers_context.bindings.get(idx);
         match binding {
             Binding::Forward(idx)
             | Binding::ForwardToFirstUse(idx)
@@ -280,8 +282,8 @@ impl<'a> CapturedVariableVisitor<'a> {
 impl<'a> AstScopedVisitor for CapturedVariableVisitor<'a> {
     fn on_scope_update(&mut self, scopes: &Scopes) {
         self.current_exported_function = scopes.current_exported_function(
-            self.module_context.module_id,
-            self.module_context.module_info.name(),
+            self.module_context.answers_context.module_id,
+            self.module_context.answers_context.module_info.name(),
             &SCOPE_EXPORTED_FUNCTION_FLAGS,
         );
     }
@@ -387,7 +389,7 @@ pub fn export_captured_variables_for_module(
     context: &ModuleContext,
 ) -> HashMap<FunctionRef, Vec<CapturedVariableRef<FunctionRef>>> {
     captured_variables
-        .get_for_module(context.module_id)
+        .get_for_module(context.answers_context.module_id)
         .unwrap()
         .clone()
         .0

pyrefly/lib/report/pysa/captured_variable.rs

@@ -44,6 +44,7 @@ use crate::report::pysa::ModuleContext;
 use crate::report::pysa::call_graph::Target;
 use crate::report::pysa::call_graph::resolve_decorator_callees;
 use crate::report::pysa::collect::CollectNoDuplicateKeys;
+use crate::report::pysa::context::ModuleAnswersContext;
 use crate::report::pysa::function::FunctionBaseDefinition;
 use crate::report::pysa::function::FunctionRef;
 use crate::report::pysa::function::WholeProgramFunctionDefinitions;
@@ -232,7 +233,7 @@ impl ClassDefinition {
     }
 }
 
-pub fn get_all_classes(context: &ModuleContext) -> impl Iterator<Item = Class> {
+pub fn get_all_classes(context: &ModuleAnswersContext) -> impl Iterator<Item = Class> {
     context
         .bindings
         .keys::<KeyClass>()
@@ -242,14 +243,29 @@ pub fn get_all_classes(context: &ModuleContext) -> impl Iterator<Item = Class> {
 pub fn get_class_field_from_current_class_only(
     class: &Class,
     field_name: &Name,
-    context: &ModuleContext,
+    context: &ModuleAnswersContext,
 ) -> Option<Arc<ClassField>> {
-    context
-        .transaction
-        .ad_hoc_solve(&context.handle, "pysa_class_field", |solver| {
-            solver.get_field_from_current_class_only(class, field_name)
-        })
-        .unwrap()
+    // This inlines the logic from `AnswersSolver::get_field_from_current_class_only`,
+    // `get_non_synthesized_field_from_current_class_only`, and
+    // `get_synthesized_field_from_current_class_only`.
+    assert!(class.module() == &context.module_info);
+
+    // Non-synthesized field: check class fields list, then look up the answer.
+    let class_fields = &context.bindings.metadata().get_class(class.index()).fields;
+    if class_fields.contains(field_name) {
+        let key = KeyClassField(class.index(), field_name.clone());
+        if let Some(idx) = context.bindings.key_to_idx_hashed_opt(Hashed::new(&key))
+            && let Some(field) = context.answers.get_idx(idx)
+        {
+            return Some(field);
+        }
+    }
+
+    // Synthesized field (e.g., dataclass fields).
+    let key = KeyClassSynthesizedFields(class.index());
+    let idx = context.bindings.key_to_idx_hashed_opt(Hashed::new(&key))?;
+    let synthesized_fields = context.answers.get_idx(idx)?;
+    Some(synthesized_fields.get(field_name)?.inner.dupe())
 }
 
 pub fn get_super_class_member(
@@ -260,16 +276,18 @@ pub fn get_super_class_member(
 ) -> Option<WithDefiningClass<Arc<ClassField>>> {
     context
         .transaction
-        .ad_hoc_solve(&context.handle, "pysa_super_class_member", |solver| {
-            solver.get_super_class_member(class, start_lookup_cls, field_name)
-        })
+        .ad_hoc_solve(
+            &context.answers_context.handle,
+            "pysa_super_class_member",
+            |solver| solver.get_super_class_member(class, start_lookup_cls, field_name),
+        )
         .flatten()
 }
 
 pub fn get_class_field_declaration<'a>(
     class: &Class,
     field_name: &Name,
-    context: &'a ModuleContext,
+    context: &'a ModuleAnswersContext,
 ) -> Option<&'a BindingClassField> {
     assert_eq!(class.module(), &context.module_info);
     let key_class_field = KeyClassField(class.index(), field_name.clone());
@@ -280,7 +298,7 @@ pub fn get_class_field_declaration<'a>(
         .map(|idx| context.bindings.get(idx))
 }
 
-pub fn get_class_mro(class: &Class, context: &ModuleContext) -> Arc<ClassMro> {
+pub fn get_class_mro(class: &Class, context: &ModuleAnswersContext) -> Arc<ClassMro> {
     assert_eq!(class.module(), &context.module_info);
     context
         .answers
@@ -290,7 +308,7 @@ pub fn get_class_mro(class: &Class, context: &ModuleContext) -> Arc<ClassMro> {
 
 pub fn get_class_fields<'a>(
     class: &'a Class,
-    context: &'a ModuleContext<'a>,
+    context: &'a ModuleAnswersContext,
 ) -> impl Iterator<Item = (Cow<'a, Name>, Arc<ClassField>)> {
     let class_fields = context
         .bindings
@@ -368,11 +386,11 @@ fn export_class_fields(
     context: &ModuleContext,
     ann_assign_map: &AnnAssignMap,
 ) -> HashMap<Name, PysaClassField> {
-    assert_eq!(class.module(), &context.module_info);
-    get_class_fields(class, context)
+    assert_eq!(class.module(), &context.answers_context.module_info);
+    get_class_fields(class, &context.answers_context)
         .filter(|(_, field)| !is_callable_like(&field.ty()))
         .filter_map(|(name, field)| {
-            let field_binding = get_class_field_declaration(class, &name, context);
+            let field_binding = get_class_field_declaration(class, &name, &context.answers_context);
 
             let explicit_annotation = match field_binding {
                 Some(BindingClassField {
@@ -389,7 +407,7 @@ fn export_class_fields(
                 }) => *annotation,
                 _ => None,
             }
-            .map(|idx| context.bindings.idx_to_key(idx))
+            .map(|idx| context.answers_context.bindings.idx_to_key(idx))
             .and_then(|key_annotation| match key_annotation {
                 // We want to export the annotation as it is in the source code.
                 // We cannot use the answer for `key_annotation` (which wraps a `Type`),
@@ -398,11 +416,19 @@ fn export_class_fields(
                 KeyAnnotation::Annotation(identifier) => ann_assign_map
                     .get(identifier.range().start())
                     .map(|annotation_range| {
-                        context.module_info.code_at(*annotation_range).to_owned()
+                        context
+                            .answers_context
+                            .module_info
+                            .code_at(*annotation_range)
+                            .to_owned()
                     }),
-                KeyAnnotation::AttrAnnotation(range) => {
-                    Some(context.module_info.code_at(*range).to_owned())
-                }
+                KeyAnnotation::AttrAnnotation(range) => Some(
+                    context
+                        .answers_context
+                        .module_info
+                        .code_at(*range)
+                        .to_owned(),
+                ),
                 _ => None,
             });
 
@@ -421,7 +447,10 @@ fn export_class_fields(
                     PysaClassField {
                         type_: PysaType::from_type(&field.ty(), context),
                         explicit_annotation,
-                        location: Some(PysaLocation::from_text_range(*range, &context.module_info)),
+                        location: Some(PysaLocation::from_text_range(
+                            *range,
+                            &context.answers_context.module_info,
+                        )),
                         declaration_kind: PysaClassFieldDeclaration::from(definition),
                     },
                 )),
@@ -444,8 +473,8 @@ fn find_definition_ast<'a>(
     class: &Class,
     context: &'a ModuleContext<'a>,
 ) -> Option<&'a StmtClassDef> {
-    assert_eq!(class.module(), &context.module_info);
-    Ast::locate_node(&context.ast, class.qname().range().start())
+    assert_eq!(class.module(), &context.answers_context.module_info);
+    Ast::locate_node(&context.answers_context.ast, class.qname().range().start())
         .iter()
         .find_map(|node| match node {
             AnyNodeRef::StmtClassDef(stmt) if stmt.name.range == class.qname().range() => {
@@ -461,7 +490,7 @@ fn get_decorator_callees(
     function_base_definitions: &WholeProgramFunctionDefinitions<FunctionBaseDefinition>,
     context: &ModuleContext,
 ) -> HashMap<PysaLocation, Vec<Target<FunctionRef>>> {
-    assert_eq!(class.module(), &context.module_info);
+    assert_eq!(class.module(), &context.answers_context.module_info);
     if let Some(class_def) = find_definition_ast(class, context) {
         resolve_decorator_callees(
             &class_def.decorator_list,
@@ -480,24 +509,35 @@ pub fn export_all_classes(
     context: &ModuleContext,
 ) -> HashMap<PysaLocation, ClassDefinition> {
     let mut class_definitions = HashMap::new();
-    let ann_assign_map = AnnAssignMap::build(&context.ast);
+    let ann_assign_map = AnnAssignMap::build(&context.answers_context.ast);
 
-    for class_idx in context.bindings.keys::<KeyClass>() {
+    for class_idx in context.answers_context.bindings.keys::<KeyClass>() {
         let class = context
+            .answers_context
             .answers
             .get_idx(class_idx)
             .unwrap()
             .0
             .dupe()
             .unwrap();
         let class_index = class.index();
-        let parent = get_scope_parent(&context.ast, &context.module_info, class.qname().range());
+        let parent = get_scope_parent(
+            &context.answers_context.ast,
+            &context.answers_context.module_info,
+            class.qname().range(),
+        );
         let metadata = context
+            .answers_context
             .answers
-            .get_idx(context.bindings.key_to_idx(&KeyClassMetadata(class_index)))
+            .get_idx(
+                context
+                    .answers_context
+                    .bindings
+                    .key_to_idx(&KeyClassMetadata(class_index)),
+            )
             .unwrap();
 
-        let is_synthesized = match context.bindings.get(class_idx) {
+        let is_synthesized = match context.answers_context.bindings.get(class_idx) {
             BindingClass::FunctionalClassDef(_, _, _) => true,
             BindingClass::ClassDef(_) => false,
         };
@@ -510,7 +550,7 @@ pub fn export_all_classes(
             .map(|base_class| ClassRef::from_class(base_class, context.module_ids))
             .collect::<Vec<_>>();
 
-        let mro = match &*get_class_mro(&class, context) {
+        let mro = match &*get_class_mro(&class, &context.answers_context) {
             ClassMro::Resolved(mro) => PysaClassMro::Resolved(
                 mro.iter()
                     .map(|class_type| {
@@ -545,7 +585,10 @@ pub fn export_all_classes(
         assert!(
             class_definitions
                 .insert(
-                    PysaLocation::from_text_range(class.qname().range(), &context.module_info),
+                    PysaLocation::from_text_range(
+                        class.qname().range(),
+                        &context.answers_context.module_info
+                    ),
                     class_definition
                 )
                 .is_none(),