feat: self state across methods

README.md

@@ -129,6 +129,12 @@ allocation site one abstract object with an `elements` and an `attributes` locat
 `insert`, `add`, `update` and iteration all carry taint through the container. Function
 summaries record the parameters a function mutates and the module globals it touches, so
 a helper that fills a list taints the caller's list, across files too.
+Objects of the project's own classes are followed too: `App(x)` calls `App.__init__` on
+the new object, `app.run()` and `self.run()` call the class's methods with the receiver as
+`self`, so a value stored in one method and read in another flows through the object,
+across files through the project index. A method the framework calls, a Django view's
+`get` for instance, starts with what its sibling methods store into `self` from their
+own inputs.
 Each flow is then judged: a dominating guard that proves the value safe (`isdigit()`,
 membership in a constant allowlist, equality with a constant) refutes it and it is not
 reported; a guard that only mentions the value makes it a hotspot with medium confidence;

src/coretrace_python/interprocedural/callgraph.py

@@ -151,29 +151,48 @@ def resolve_targets(
     known: frozenset[str],
     parameters: Mapping[Value, SymbolId] | None = None,
     nested: frozenset[str] = frozenset(),
+    classes: Mapping[str, frozenset[str]] | None = None,
+    owner: str | None = None,
+    typed: Mapping[Value, str] | None = None,
 ) -> tuple[dict[Value, Target], dict[Value, SymbolId]]:
     """Map every callee value of ``function`` to its target, and every symbol value.
-    ``nested`` names the functions defined inside this one, ``outer.inner``."""
+    ``nested`` names the functions defined inside this one (``outer.inner``);
+    ``classes`` maps the module's classes to their methods, ``owner`` is the class of a
+    method and ``typed`` the parameters annotated with a module class. ``App(x)`` is a
+    call to ``App.__init__``, ``app.run()`` and ``self.run()`` calls to ``App.run``."""
 
     module = scopes.module_scope
     symbols = derive_symbols(function, parameters)
     targets: dict[Value, Target] = {
         value: ExternalSymbol(symbol) for value, symbol in symbols.items()
     }
+    classes = classes or {}
+    instance_of: dict[Value, str] = dict(typed or {})
+    if owner is not None and function.parameters:
+        instance_of[function.parameters[0]] = owner
+    defs = {i.result: i for block in function.blocks for i in block.instructions if i.result}
     for block in function.blocks:
         for instruction in block.instructions:
             if isinstance(instruction, MakeFunction):
                 qualified = f"{function.name}.{instruction.name}"
                 if qualified in nested:
                     targets[instruction.result] = KnownFunction(qualified)
-            if isinstance(instruction, Global):
+            elif isinstance(instruction, Global):
                 binding = module.bindings.get(instruction.name)
-                if (
-                    binding is not None
-                    and binding.kind is BindingKind.FUNCTION
-                    and instruction.name in known
-                ):
+                if binding is None:
+                    continue
+                if binding.kind is BindingKind.FUNCTION and instruction.name in known:
                     targets[instruction.result] = KnownFunction(instruction.name)
+                elif binding.kind is BindingKind.CLASS and "__init__" in classes.get(instruction.name, ()):
+                    targets[instruction.result] = KnownFunction(f"{instruction.name}.__init__")
+            elif isinstance(instruction, Call):
+                callee = defs.get(instruction.callee)
+                if isinstance(callee, Global) and callee.name in classes:
+                    instance_of[instruction.result] = callee.name
+            elif isinstance(instruction, GetAttr) and instruction.object in instance_of:
+                method = f"{instance_of[instruction.object]}.{instruction.attribute}"
+                if method in nested:
+                    targets[instruction.result] = KnownFunction(method)
     return targets, symbols
 
 
@@ -195,6 +214,19 @@ def _annotated(
     return found
 
 
+def _typed_with_module_classes(
+    function: nodes.Function, ssa: FunctionIR, classes: Mapping[str, frozenset[str]]
+) -> dict[Value, str]:
+    """Parameters annotated with a class of this module (``a: App``)."""
+
+    found: dict[Value, str] = {}
+    for value, parameter in zip(ssa.parameters, function.parameters, strict=False):
+        annotation = parameter.annotation
+        if isinstance(annotation, nodes.Name) and annotation.identifier in classes:
+            found[value] = annotation.identifier
+    return found
+
+
 class CallGraphAnalysis(Analysis[CallGraph]):
     name: ClassVar[str] = "interprocedural.callgraph"
     requires: ClassVar[frozenset[AnyAnalysis]] = frozenset({SSAAnalysis, ScopeAnalysis, SymbolAnalysis})
@@ -216,6 +248,11 @@ def compute(cls, ctx: AnalysisContext) -> CallGraph:
         known = frozenset(
             name for name, function in definitions.items() if function in ctx.module.body
         )
+        classes = {
+            statement.name: frozenset(m.name for m in statement.body if isinstance(m, nodes.Function))
+            for statement in ctx.module.body
+            if isinstance(statement, nodes.Class)
+        }
         sites: dict[str, tuple[CallSite, ...]] = {}
         symbols: dict[str, Mapping[Value, SymbolId]] = {}
         unsupported: set[str] = set()
@@ -226,8 +263,16 @@ def compute(cls, ctx: AnalysisContext) -> CallGraph:
                 unsupported.add(name)
                 sites[name] = ()
                 continue
+            owner = name.rsplit(".", 1)[0] if "." in name and name.rsplit(".", 1)[0] in classes else None
             targets, symbols[name] = resolve_targets(
-                ssa, scopes, known, _annotated(function, ssa, scopes, table), frozenset(definitions)
+                ssa,
+                scopes,
+                known,
+                _annotated(function, ssa, scopes, table),
+                frozenset(definitions),
+                classes,
+                owner,
+                _typed_with_module_classes(function, ssa, classes),
             )
             found: list[CallSite] = []
             for block in ssa.blocks:

src/coretrace_python/interprocedural/summaries.py

@@ -300,8 +300,12 @@ def call(self, call: Call, state: dict[Key, Dep]) -> Dep:
 
         if isinstance(target, ExternalSymbol):
             project = self.project.get(target.symbol)
+            if project is None:
+                project = self.project.get(target.symbol.attribute("__init__"))
             if project is not None:
-                return self.known(project, arguments, keywords, everything, call, state)
+                bound = self.receiver(call, project.name)
+                arguments = (*(self.deep(r, state) for r in bound), *arguments)
+                return self.known(project, arguments, keywords, everything, call, state, bound)
             self.record(
                 target.symbol,
                 tuple(a.parameters for a in arguments),
@@ -314,10 +318,25 @@ def call(self, call: Call, state: dict[Key, Dep]) -> Dep:
             made = self.defs.get(call.callee)
             if isinstance(made, MakeFunction):
                 arguments = (*arguments, *(self.deep(v, state) for v in made.captured))
-            return self.known(self.table[target.name], arguments, keywords, everything, call, state)
+            bound = self.receiver(call, target.name)
+            arguments = (*(self.deep(r, state) for r in bound), *arguments)
+            return self.known(self.table[target.name], arguments, keywords, everything, call, state, bound)
         assert isinstance(target, UnknownTarget)
         return everything | state.get(call.callee, EMPTY)
 
+    def receiver(self, call: Call, name: str) -> tuple[Value, ...]:
+        """The object a method call runs on: the receiver of ``obj.method(...)`` or the
+        new object of ``Class(...)``, the method's implicit first parameter."""
+
+        if "." not in name:
+            return ()
+        callee = self.defs.get(call.callee)
+        if isinstance(callee, GetAttr):
+            return (callee.object,)
+        if name.endswith(".__init__") and isinstance(callee, Global | Symbol):
+            return (call.result,)
+        return ()
+
     def known(
         self,
         callee: FunctionSummary,
@@ -326,6 +345,7 @@ def known(
         everything: Dep,
         call: Call,
         state: dict[Key, Dep],
+        receiver: tuple[Value, ...] = (),
     ) -> Dep:
         """Dependencies of a call to a function whose summary is known."""
 
@@ -347,7 +367,7 @@ def mapped(indices: Dependencies) -> Dep:
                 call.location,
             )
         made = self.defs.get(call.callee)
-        values = (*call.arguments, *(made.captured if isinstance(made, MakeFunction) else ()))
+        values = (*receiver, *call.arguments, *(made.captured if isinstance(made, MakeFunction) else ()))
         for mutation in callee.mutations:
             if mutation.parameter < len(values):
                 deps = mapped(mutation.dependencies) | Dep(externals=mutation.externals)