Export decorated target call graphs

Summary:
In pysa, we create artificial functions called "decorated targets" that correspond to the application of decorators on a function.

For instance:
```
foo(1)
bar
def baz(): ...
```

`baz` will have a decorated targets:
```
def baz@decorated():
  return foo(1)(bar(baz))
```

We need the call graph of this function from pyrefly.
This diff implements that logic: when computing call graphs, expressions in decorators are exported in a special `FunctionDecoratedTarget` function id.

Reviewed By: tianhan0

Differential Revision: D85569064

fbshipit-source-id: e5f27d1c63fae8136c6110da4a775e42c0063580

Author: arthaud

pyrefly/lib/report/pysa/ast_visitor.rs

@@ -76,84 +76,147 @@ pub struct Scopes {
     stack: Vec<Scope>,
 }
 
+pub enum ExportFunctionDecorators {
+    InDecoratedFunction,
+    InParentScope,
+    InDecoratedTarget,
+    Ignore,
+}
+
+pub enum ExportClassDecorators {
+    InDecoratedClassTopLevel,
+    InParentScope,
+    Ignore,
+}
+
+pub enum ExportDefaultArguments {
+    InFunction,
+    InParentScope,
+    Ignore,
+}
+
+pub struct ScopeExportedFunctionFlags {
+    pub include_top_level: bool,
+    pub include_class_top_level: bool,
+    pub include_function_decorators: ExportFunctionDecorators,
+    pub include_class_decorators: ExportClassDecorators,
+    pub include_default_arguments: ExportDefaultArguments,
+}
+
 impl Scopes {
     pub fn current_exported_function(
         &self,
         module_id: ModuleId,
         module_name: ModuleName,
-        include_top_level: bool,
-        include_class_top_level: bool,
-        include_decorators_in_decorated_definition: bool,
-        include_default_arguments_in_function: bool,
+        flags: ScopeExportedFunctionFlags,
     ) -> Option<FunctionRef> {
-        let mut iterator = self.stack.iter().rev();
-        loop {
-            match iterator.next().unwrap() {
-                Scope::TopLevel => {
-                    if include_top_level {
-                        return Some(FunctionRef {
-                            module_id,
-                            module_name,
-                            function_id: FunctionId::ModuleTopLevel,
-                            function_name: Name::from("$toplevel"),
-                        });
-                    } else {
-                        return None;
-                    }
+        Self::current_exported_function_impl(self.stack.iter().rev(), module_id, module_name, flags)
+    }
+
+    fn current_exported_function_impl<'a>(
+        mut iterator: impl Iterator<Item = &'a Scope>,
+        module_id: ModuleId,
+        module_name: ModuleName,
+        flags: ScopeExportedFunctionFlags,
+    ) -> Option<FunctionRef> {
+        match iterator.next().unwrap() {
+            Scope::TopLevel => {
+                if flags.include_top_level {
+                    Some(FunctionRef {
+                        module_id,
+                        module_name,
+                        function_id: FunctionId::ModuleTopLevel,
+                        function_name: Name::from("$toplevel"),
+                    })
+                } else {
+                    None
                 }
-                Scope::ExportedFunction {
-                    function_id,
-                    function_name,
-                    ..
-                } => {
-                    return Some(FunctionRef {
+            }
+            Scope::ExportedFunction {
+                function_id,
+                function_name,
+                ..
+            } => Some(FunctionRef {
+                module_id,
+                module_name,
+                function_id: function_id.clone(),
+                function_name: function_name.clone(),
+            }),
+            Scope::NonExportedFunction { .. } => None,
+            Scope::ExportedClass { class_id, .. } => {
+                if flags.include_class_top_level {
+                    Some(FunctionRef {
                         module_id,
                         module_name,
-                        function_id: function_id.clone(),
-                        function_name: function_name.clone(),
-                    });
+                        function_id: FunctionId::ClassTopLevel {
+                            class_id: *class_id,
+                        },
+                        function_name: Name::from("$class_toplevel"),
+                    })
+                } else {
+                    None
                 }
-                Scope::NonExportedFunction { .. } => {
-                    return None;
+            }
+            Scope::NonExportedClass { .. } => None,
+            Scope::FunctionParameters => match flags.include_default_arguments {
+                ExportDefaultArguments::InFunction => {
+                    Self::current_exported_function_impl(iterator, module_id, module_name, flags)
                 }
-                Scope::ExportedClass { class_id, .. } => {
-                    if include_class_top_level {
-                        return Some(FunctionRef {
-                            module_id,
-                            module_name,
-                            function_id: FunctionId::ClassTopLevel {
-                                class_id: *class_id,
-                            },
-                            function_name: Name::from("$class_toplevel"),
-                        });
-                    } else {
-                        return None;
-                    }
+                ExportDefaultArguments::InParentScope => {
+                    iterator.next().unwrap();
+                    Self::current_exported_function_impl(iterator, module_id, module_name, flags)
                 }
-                Scope::NonExportedClass { .. } => {
-                    return None;
+                ExportDefaultArguments::Ignore => None,
+            },
+            Scope::FunctionTypeParams
+            | Scope::FunctionReturnAnnotation
+            | Scope::ClassTypeParams
+            | Scope::ClassArguments => {
+                // These are not true "semantic" scopes.
+                // We need to skip the parent scope, which is the wrapping function/class scope.
+                iterator.next().unwrap();
+                Self::current_exported_function_impl(iterator, module_id, module_name, flags)
+            }
+            Scope::ClassDecorators => match flags.include_class_decorators {
+                ExportClassDecorators::InDecoratedClassTopLevel => {
+                    Self::current_exported_function_impl(iterator, module_id, module_name, flags)
                 }
-                Scope::FunctionParameters => {
-                    if !include_default_arguments_in_function {
-                        // This not a true "semantic" scope.
-                        // We need to skip the parent scope, which is the wrapping function/class scope.
-                        iterator.next().unwrap();
-                    }
+                ExportClassDecorators::InParentScope => {
+                    iterator.next().unwrap();
+                    Self::current_exported_function_impl(iterator, module_id, module_name, flags)
                 }
-                Scope::FunctionTypeParams
-                | Scope::FunctionReturnAnnotation
-                | Scope::ClassTypeParams
-                | Scope::ClassArguments => {
-                    // These are not true "semantic" scopes.
-                    // We need to skip the parent scope, which is the wrapping function/class scope.
+                ExportClassDecorators::Ignore => None,
+            },
+            Scope::FunctionDecorators => match flags.include_function_decorators {
+                ExportFunctionDecorators::InDecoratedFunction => {
+                    Self::current_exported_function_impl(iterator, module_id, module_name, flags)
+                }
+                ExportFunctionDecorators::InParentScope => {
                     iterator.next().unwrap();
+                    Self::current_exported_function_impl(iterator, module_id, module_name, flags)
                 }
-                Scope::FunctionDecorators | Scope::ClassDecorators => {
-                    if !include_decorators_in_decorated_definition {
-                        iterator.next().unwrap();
+                ExportFunctionDecorators::InDecoratedTarget => {
+                    match Self::current_exported_function_impl(
+                        iterator,
+                        module_id,
+                        module_name,
+                        flags,
+                    ) {
+                        Some(FunctionRef {
+                            function_id: FunctionId::Function { location },
+                            function_name,
+                            ..
+                        }) => Some(FunctionRef {
+                            module_id,
+                            module_name,
+                            function_id: FunctionId::FunctionDecoratedTarget { location },
+                            function_name,
+                        }),
+                        _ => None,
                     }
                 }
-            }
+                ExportFunctionDecorators::Ignore => None,
+            },
         }
     }
 }
@@ -192,33 +255,58 @@ impl ScopeId {
 }
 
 pub trait AstScopedVisitor {
-    fn visit_statement(&mut self, _stmt: &Stmt, _scopes: &Scopes) {}
+    fn visit_statement(&mut self, stmt: &Stmt, scopes: &Scopes) {
+        let _ = stmt;
+        let _ = scopes;
+    }
     fn visit_expression(
         &mut self,
-        _expr: &Expr,
-        _scopes: &Scopes,
-        _parent_expression: Option<&Expr>,
+        expr: &Expr,
+        scopes: &Scopes,
+        parent_expression: Option<&Expr>,
         // If the current expression is in an assignment, this is the left side of the assignment
-        _assignment_targets: Option<&Vec<&Expr>>,
+        assignment_targets: Option<&Vec<&Expr>>,
     ) {
+        let _ = expr;
+        let _ = scopes;
+        let _ = parent_expression;
+        let _ = assignment_targets;
     }
     fn enter_function_scope(
         &mut self,
-        _function_def: &StmtFunctionDef,
-        _scopes_in_function: &Scopes,
+        function_def: &StmtFunctionDef,
+        scopes_in_function: &Scopes,
     ) {
+        let _ = function_def;
+        let _ = scopes_in_function;
     }
     fn exit_function_scope(
         &mut self,
-        _function_def: &StmtFunctionDef,
-        _scopes_outside_function: &Scopes,
+        function_def: &StmtFunctionDef,
+        scopes_outside_function: &Scopes,
     ) {
+        let _ = function_def;
+        let _ = scopes_outside_function;
+    }
+    fn enter_class_scope(&mut self, class_def: &StmtClassDef, scopes_in_class: &Scopes) {
+        let _ = class_def;
+        let _ = scopes_in_class;
+    }
+    fn exit_class_scope(&mut self, function_def: &StmtClassDef, scopes_outside_class: &Scopes) {
+        let _ = function_def;
+        let _ = scopes_outside_class;
+    }
+    fn enter_toplevel_scope(&mut self, ast: &ModModule, scopes_in_toplevel: &Scopes) {
+        let _ = ast;
+        let _ = scopes_in_toplevel;
+    }
+    fn exit_toplevel_scope(&mut self, ast: &ModModule, scopes_in_toplevel: &Scopes) {
+        let _ = ast;
+        let _ = scopes_in_toplevel;
+    }
+    fn on_scope_update(&mut self, scopes: &Scopes) {
+        let _ = scopes;
     }
-    fn enter_class_scope(&mut self, _class_def: &StmtClassDef, _scopes_in_class: &Scopes) {}
-    fn exit_class_scope(&mut self, _function_def: &StmtClassDef, _scopes_outside_class: &Scopes) {}
-    fn enter_toplevel_scope(&mut self, _ast: &ModModule, _scopes_in_toplevel: &Scopes) {}
-    fn exit_toplevel_scope(&mut self, _ast: &ModModule, _scopes_in_toplevel: &Scopes) {}
-    fn on_scope_update(&mut self, _scopes: &Scopes) {}
     fn visit_type_annotations() -> bool;
 }
 
@@ -305,26 +393,28 @@ fn visit_statement<V: AstScopedVisitor>(
             visitor.enter_function_scope(function_def, scopes);
             visitor.on_scope_update(scopes);
 
-            scopes.stack.push(Scope::FunctionDecorators);
-            visitor.on_scope_update(scopes);
-            function_def.decorator_list.iter().for_each(&mut |e| {
-                ruff_python_ast::visitor::source_order::SourceOrderVisitor::visit_decorator(
-                    &mut CustomSourceOrderVisitor {
-                        visitor,
-                        scopes,
-                        module_context,
-                        parent_expression: None,
-                        assignment_targets: None,
-                    },
-                    e,
-                )
-            });
-            scopes.stack.pop();
-            visitor.on_scope_update(scopes);
+            if !function_def.decorator_list.is_empty() {
+                scopes.stack.push(Scope::FunctionDecorators);
+                visitor.on_scope_update(scopes);
+                function_def.decorator_list.iter().for_each(&mut |e| {
+                    ruff_python_ast::visitor::source_order::SourceOrderVisitor::visit_decorator(
+                        &mut CustomSourceOrderVisitor {
+                            visitor,
+                            scopes,
+                            module_context,
+                            parent_expression: None,
+                            assignment_targets: None,
+                        },
+                        e,
+                    )
+                });
+                scopes.stack.pop();
+                visitor.on_scope_update(scopes);
+            }
 
-            scopes.stack.push(Scope::FunctionTypeParams);
-            visitor.on_scope_update(scopes);
             if let Some(type_params) = &function_def.type_params {
+                scopes.stack.push(Scope::FunctionTypeParams);
+                visitor.on_scope_update(scopes);
                 ruff_python_ast::visitor::source_order::SourceOrderVisitor::visit_type_params(
                     &mut CustomSourceOrderVisitor {
                         visitor,
@@ -335,9 +425,9 @@ fn visit_statement<V: AstScopedVisitor>(
                     },
                     type_params,
                 );
+                scopes.stack.pop();
+                visitor.on_scope_update(scopes);
             }
-            scopes.stack.pop();
-            visitor.on_scope_update(scopes);
 
             scopes.stack.push(Scope::FunctionParameters);
             visitor.on_scope_update(scopes);
@@ -354,22 +444,24 @@ fn visit_statement<V: AstScopedVisitor>(
             scopes.stack.pop();
             visitor.on_scope_update(scopes);
 
-            scopes.stack.push(Scope::FunctionReturnAnnotation);
-            visitor.on_scope_update(scopes);
-            function_def.returns.iter().for_each(|return_annotation| {
-                ruff_python_ast::visitor::source_order::SourceOrderVisitor::visit_annotation(
-                    &mut CustomSourceOrderVisitor {
-                        visitor,
-                        scopes,
-                        module_context,
-                        parent_expression: None,
-                        assignment_targets: None,
-                    },
-                    return_annotation,
-                );
-            });
-            scopes.stack.pop();
-            visitor.on_scope_update(scopes);
+            if function_def.returns.is_some() {
+                scopes.stack.push(Scope::FunctionReturnAnnotation);
+                visitor.on_scope_update(scopes);
+                function_def.returns.iter().for_each(|return_annotation| {
+                    ruff_python_ast::visitor::source_order::SourceOrderVisitor::visit_annotation(
+                        &mut CustomSourceOrderVisitor {
+                            visitor,
+                            scopes,
+                            module_context,
+                            parent_expression: None,
+                            assignment_targets: None,
+                        },
+                        return_annotation,
+                    );
+                });
+                scopes.stack.pop();
+                visitor.on_scope_update(scopes);
+            }
 
             for stmt in &function_def.body {
                 visit_statement(stmt, visitor, scopes, module_context);