Follow re-export chains when checking call safety

Re-exported symbols (e.g. bar.Foo re-exporting foo.Foo) were not being
resolved to their original definitions during safety analysis. This
caused check_call to miss unsafe constructors/functions accessed through
re-exports, incorrectly marking them as safe.

Change re_exports from a HashSet to a HashMap that preserves the mapping
to original names, add resolve_re_export() that follows chains, and
apply it in check_call and the ImportedTypeAttr property lookup.

Author: msullivan

src/project.rs

@@ -174,7 +174,10 @@ fn merge_all_functions_and_methods(
     )
 }
 
-fn get_all_safe_re_exports(effect_table: &EffectTable, re_exports: &mut AHashSet<ModuleName>) {
+fn get_all_safe_re_exports(
+    effect_table: &EffectTable,
+    re_exports: &mut AHashMap<ModuleName, ModuleName>,
+) {
     let unsafe_re_exports = effect_table
         .values()
         .flatten()
@@ -790,7 +793,7 @@ struct ProjectInfo {
     classes: ClassTable,
     // Mappings of functions to the containing module
     functions: AHashMap<ModuleName, ModuleName>,
-    re_exports: AHashSet<ModuleName>,
+    re_exports: AHashMap<ModuleName, ModuleName>,
     // Mapping of all methods called on imported objects
     methods: AHashMap<ModuleName, ModuleName>,
     // Reverse mapping: parent function → nested function scopes.
@@ -810,9 +813,9 @@ impl ProjectInfo {
             merge_all_classes(&mut analysis_map)
         });
         let re_exports = time("    Getting re-exports", || {
-            let mut re_exports = exports
+            let mut re_exports: AHashMap<ModuleName, ModuleName> = exports
                 .get_re_exports()
-                .map(|(name, _)| name.as_module_name())
+                .map(|(name, (original, _))| (name.as_module_name(), original.as_module_name()))
                 .collect();
             get_all_safe_re_exports(&effect_table, &mut re_exports);
             re_exports
@@ -833,6 +836,22 @@ impl ProjectInfo {
         }
     }
 
+    /// Resolve a name through the re-exports table, following chains
+    /// (e.g. baz.Foo → bar.Foo → foo.Foo). Returns the original definition
+    /// name, or the name unchanged if it is not a re-export. Bails out of
+    /// cycles (e.g. `from b import X` / `from a import X`).
+    fn resolve_re_export(&self, name: &ModuleName) -> ModuleName {
+        let mut current = *name;
+        let mut seen = AHashSet::new();
+        while let Some(&original) = self.re_exports.get(&current) {
+            if !seen.insert(current) {
+                break;
+            }
+            current = original;
+        }
+        current
+    }
+
     pub fn contains_callable(&self, name: &ModuleName) -> bool {
         if self.functions.contains_key(name) || self.classes.contains(name) {
             return true;
@@ -841,7 +860,7 @@ impl ProjectInfo {
         if self.functions.contains_key(&call_name) || self.classes.contains(&call_name) {
             true
         } else {
-            self.re_exports.contains(&call_name)
+            self.re_exports.contains_key(&call_name)
         }
     }
 
@@ -943,6 +962,7 @@ impl ProjectInfo {
             } else if eff.kind == EffectKind::ImportedTypeAttr {
                 // Check if this is a property access
                 if let Some((typ, attr)) = eff.name.split_attr() {
+                    let typ = self.resolve_re_export(&typ);
                     if let Some(field) = self
                         .classes
                         .lookup(&typ)
@@ -1024,11 +1044,15 @@ impl ProjectInfo {
     }
 
     fn check_call(&self, call: &mut Call, state: &GlobalAnalysisState) -> Result<bool> {
-        if self.classes.contains(&call.func) {
+        // Resolve re-exports to their original name so that class/function
+        // lookups find the actual definition (e.g. bar.Foo → foo.Foo).
+        let resolved_func = self.resolve_re_export(&call.func);
+        let mut resolved_call = call.clone_with_name(resolved_func);
+        if self.classes.contains(&resolved_call.func) {
             // This is a class constructor
-            self.check_constructor_call(call, state)
+            self.check_constructor_call(&resolved_call, state)
         } else {
-            self.check_call_body(call, state)
+            self.check_call_body(&mut resolved_call, state)
         }
     }
 

tests/calls.rs

@@ -222,8 +222,7 @@ import importlib
 from importlib import import_module
 
 a = importlib.import_module("sys")
-# This depends on a chain of import aliases which we don't handle well
-b = importlib.__import__("math") # TODO: unsafe-function-call
+b = importlib.__import__("math") # E: unsafe-function-call
 
 import_module("bar")
 "#;
@@ -478,4 +477,91 @@ foo.bar.f()
 "#;
         check_all(vec![("foo.bar", foo_bar), ("__main__", __main__)]);
     }
+
+    #[test]
+    fn test_reexported_unsafe_class() {
+        let foo = r#"
+class Foo:
+    def __init__(self) -> None:
+        raise Exception
+"#;
+        let bar = r#"
+from foo import Foo
+"#;
+        let baz = r#"
+from bar import Foo
+x = Foo()  # E: unsafe-function-call
+"#;
+        check_all(vec![("foo", foo), ("bar", bar), ("baz", baz)]);
+    }
+
+    #[test]
+    fn test_reexported_safe_class() {
+        let foo = r#"
+class Foo:
+    def __init__(self) -> None:
+        pass
+"#;
+        let bar = r#"
+from foo import Foo
+"#;
+        let baz = r#"
+from bar import Foo
+x = Foo()
+"#;
+        check_all(vec![("foo", foo), ("bar", bar), ("baz", baz)]);
+    }
+
+    #[test]
+    fn test_reexported_chain_unsafe_class() {
+        let foo = r#"
+class Foo:
+    def __init__(self) -> None:
+        raise Exception
+"#;
+        let bar = r#"
+from foo import Foo
+"#;
+        let baz = r#"
+from bar import Foo
+"#;
+        let consumer = r#"
+from baz import Foo
+x = Foo()  # E: unsafe-function-call
+"#;
+        check_all(vec![("foo", foo), ("bar", bar), ("baz", baz), ("consumer", consumer)]);
+    }
+
+    #[test]
+    fn test_reexport_cycle_terminates() {
+        // `a.X` and `b.X` re-export each other with no real definition.
+        // resolve_re_export must break out of the cycle rather than spin.
+        let a = r#"
+from b import X
+"#;
+        let b = r#"
+from a import X
+"#;
+        let main = r#"
+from a import X
+x = X()
+"#;
+        check_all(vec![("a", a), ("b", b), ("main", main)]);
+    }
+
+    #[test]
+    fn test_reexported_unsafe_function() {
+        let foo = r#"
+def f():
+    raise Exception
+"#;
+        let bar = r#"
+from foo import f
+"#;
+        let baz = r#"
+from bar import f
+f()  # E: unsafe-function-call
+"#;
+        check_all(vec![("foo", foo), ("bar", bar), ("baz", baz)]);
+    }
 }