native: annotate local functions; a real program through mypy

DESIGN.md

@@ -433,16 +433,22 @@ The ladder is also faster on CPython by itself, since a class pattern goes throu
 and positional binding at run time: on `bench/expr_eval` (ADT allocation and matching) the output
 goes from 2.30x to 1.63x of the hand-written baseline on CPython 3.14, before any compilation.
 
-**Annotations.** Top-level functions carry annotations taken from their inferred types where Python
-can spell them: `int`, `float`, `bool`, `str`, `None`, `list`/`set`/`dict`/`tuple` of those, a
+**Annotations.** Top-level functions, and functions defined inside them, carry annotations taken
+from their inferred types where Python can spell them: `int`, `float`, `bool`, `str`, `None`, `list`/`set`/`dict`/`tuple` of those, a
 record's class, and a union alias per user sum type (`Expr = Num | Add | Mul | Neg`, emitted after
 its classes). A type variable, a function type or an async result stays unannotated, so the
-annotations never claim more than the checker proved. Data class fields follow suit: a field
+annotations never claim more than the checker proved, and so does a `unit` result, since Pyfun
+returns a unit call's value (`return f(x)`), which mypy rejects under `-> None`. Data class fields follow suit: a field
 declared as a user sum type or record names it (`_0: "Expr"`, a string because the alias comes
 later), and any other field without a Python class of its own is `typing.Any`, since mypyc cannot
 build a dataclass with an `object` field. The output is written for mypy's newer redefinition rule
 (`--allow-redefinition-new --local-partial-types`): a Python local reused across match arms can hold
-differently narrowed types in different arms, which the default first-assignment rule rejects.
+differently narrowed types in different arms, which the default first-assignment rule rejects. A
+build also turns off three of mypy's error codes that only restate the Python boundary:
+`func-returns-value` (Pyfun binds a unit call's result, `_ = f(x)`), and `var-annotated` and
+`union-attr` (a module-level value mypy cannot infer, a typeshed union at an `extern` call), where
+the `extern` declarations are the contract and Pyfun's checker has already held the program to
+them.
 
 **Building.** `pyfun build --native <file> -o <dir>` does the whole trip: it compiles natively, runs
 mypyc with those flags, and leaves the extension and a `__main__.py` in `<dir>`, so the program runs

ROADMAP.md

@@ -511,9 +511,14 @@ map_build 1.64× vs hand-written.
   3.12: mypyc runs the benchmarks 1.6x / **13x** / 1.3x faster than the same program as plain Python
   (expr_eval / collatz / map_build), at **1.37x / 0.08x / 1.32x of hand-written Python**, so the
   int-arithmetic benchmark beats the hand-written baseline about 12x. On plain CPython the ladders
-  alone take expr_eval from 2.30x to 1.63x. **Remaining:** annotations for block-local functions,
-  or-patterns whose alternatives bind names (the one match shape still left to `match`), trying a
-  mypyc build of a real program (the Scrabble game), and the closure/`seq` audit under mypyc.
+  alone take expr_eval from 2.30x to 1.63x. Block-local functions are annotated too.
+  **The real-program audit (the Scrabble game, locally):** its native output passes plain mypy with
+  the build's flags (64 errors down to 0 after the shared runtime's fields became `Any`, unit
+  returns went bare, and three boundary error codes were disabled), and every replay prints the same.
+  mypyc itself still stops on **"Cannot determine type" for module-level values another module reads
+  at import time** (nullary singletons like `board._QW`, layout tables): that is the next native
+  blocker. **Remaining:** that, or-patterns whose alternatives bind names, and the closure/`seq`
+  audit under mypyc.
   The original reasoning, kept for the record:
   the checker knows every binding's inferred type, so the emitter could produce fully annotated
   Python whose annotations cannot lie, then compile it with mypyc into a C extension — native speed

src/lowering/mod.rs

@@ -136,6 +136,9 @@ pub fn lower_native_typed(
     let mut lowerer = Lowerer::new(module);
     lowerer.native = native;
     lowerer.binding_types = binding_types;
+    if native {
+        lowerer.native_spell = native_type_spellings(module);
+    }
     lowerer.float_literals = float_literals.clone();
     lowerer.order = order;
     lowerer.codecs = codecs.clone();
@@ -209,6 +212,9 @@ pub fn lower_in_project(
     let mut lowerer = Lowerer::new(module);
     lowerer.native = native;
     lowerer.binding_types = binding_types;
+    if native {
+        lowerer.native_spell = native_type_spellings(module);
+    }
     lowerer.float_literals = float_literals.clone();
     lowerer.codecs = codecs.clone();
     lowerer.imported_modules = ctx.modules.clone();
@@ -265,6 +271,13 @@ pub fn runtime_module() -> PyModule {
     PyModule { body }
 }
 
+/// [`runtime_module`] for a native build: its data class fields say `Any`
+/// rather than `object`, so a payload unwrapped from an `Option`/`Result` keeps
+/// a type mypy will pass on (`DESIGN.md` §5.6).
+pub fn runtime_module_native() -> PyModule {
+    any_for_object_fields(runtime_module())
+}
+
 /// `_pf_scope = contextvars.ContextVar("_pf_scope")`: where the innermost
 /// `Task.scope` publishes its TaskGroup for `Task.start` to find.
 fn task_scope_var() -> PyStmt {
@@ -418,6 +431,10 @@ struct Lowerer {
     /// Native mode's type table: the resolved type of each binding, keyed by the
     /// binding target's span, for the annotations on top-level functions.
     binding_types: HashMap<Span, crate::types::Ty>,
+    /// Native mode: the Python spelling of each local user type (a record's
+    /// class, a sum type's union alias), for annotating block-local functions as
+    /// they lower. See [`native_type_spellings`].
+    native_spell: HashMap<String, String>,
     /// The derived codecs the checker resolved for this module's `Decode.auto`
     /// sites (`DESIGN.md` §6, "Derived codecs").
     codecs: crate::types::Codecs,
@@ -757,6 +774,7 @@ impl Lowerer {
             needs_task_scope: false,
             native: false,
             binding_types: HashMap::new(),
+            native_spell: HashMap::new(),
             codecs: crate::types::Codecs::default(),
             codec_table_used: std::collections::BTreeMap::new(),
             needed_codec_helpers: BTreeSet::new(),
@@ -971,7 +989,7 @@ impl Lowerer {
             let returns = if *is_async {
                 None
             } else {
-                py_type_annotation(&rest, &spell)
+                return_annotation(&rest, &spell)
             };
             *stmt = PyStmt::TypedFuncDef {
                 name: name.clone(),
@@ -1833,6 +1851,23 @@ impl Lowerer {
                 self.notes.push(note);
             }
             let body = rewritten.body;
+            // Native mode: a function defined inside another is annotated here,
+            // where its binding is in hand (top-level ones are annotated by
+            // `annotate_functions`, which can also spell `Option`/`Result`).
+            if self.native
+                && !self.fn_local_stack.is_empty()
+                && let Some(ty) = self.binding_types.get(&binding.target_span.span())
+                && let Some((params, returns)) = typed_signature(&py_params, ty, &self.native_spell)
+            {
+                out.push(PyStmt::TypedFuncDef {
+                    name: py_name,
+                    params,
+                    returns,
+                    body,
+                    is_async: false,
+                });
+                return Ok(());
+            }
             out.push(PyStmt::FuncDef {
                 name: py_name,
                 params: py_params,
@@ -6533,6 +6568,67 @@ fn any_for_object_fields(mut module: PyModule) -> PyModule {
     module
 }
 
+/// The Python spelling of each non-generic local user type for native
+/// annotations: a record is its class, and a sum type its union alias (which
+/// `annotate_functions` emits after the classes), unless a constructor shares
+/// the type's name.
+fn native_type_spellings(module: &Module) -> HashMap<String, String> {
+    use crate::parser::ast::{Item, TypeDeclKind};
+    let mut spell = HashMap::new();
+    for item in &module.items {
+        if let Item::Type(decl) = item
+            && decl.params.is_empty()
+        {
+            match &decl.kind {
+                TypeDeclKind::Record(_) => {
+                    spell.insert(decl.name.clone(), py_record_class(&decl.name));
+                }
+                TypeDeclKind::Sum(variants)
+                    if !variants.is_empty() && variants.iter().all(|v| v.name != decl.name) =>
+                {
+                    spell.insert(decl.name.clone(), decl.name.clone());
+                }
+                _ => {}
+            }
+        }
+    }
+    spell
+}
+
+/// A function's annotated parameter list and return annotation from its type,
+/// peeling one arrow per parameter; `None` when the type has fewer arrows than
+/// the def has parameters.
+/// Annotated parameters (each name with its annotation, if any) and the return
+/// annotation, as [`PyStmt::TypedFuncDef`] takes them.
+type TypedSignature = (Vec<(String, Option<String>)>, Option<String>);
+
+fn typed_signature(
+    params: &[String],
+    ty: &crate::types::Ty,
+    spell: &HashMap<String, String>,
+) -> Option<TypedSignature> {
+    let mut anns = Vec::with_capacity(params.len());
+    let mut rest = ty.clone();
+    for p in params {
+        let crate::types::Ty::Fun(arg, ret, _) = rest else {
+            return None;
+        };
+        anns.push((p.clone(), py_type_annotation(&arg, spell)));
+        rest = *ret;
+    }
+    Some((anns, return_annotation(&rest, spell)))
+}
+
+/// A return annotation: as [`py_type_annotation`], except that `unit` is left
+/// unannotated, since Pyfun returns a unit call's result (`return f(x)`) and mypy
+/// rejects returning any expression from a function declared `-> None`.
+fn return_annotation(ty: &crate::types::Ty, spell: &HashMap<String, String>) -> Option<String> {
+    match ty {
+        crate::types::Ty::Unit => None,
+        other => py_type_annotation(other, spell),
+    }
+}
+
 /// A checker type spelled as a Python annotation, for native mode's typed
 /// functions (`DESIGN.md` §5.6); `None` when Python has no spelling for it here.
 /// `spell` maps a user type name to its class or union alias.

src/native_build.rs

@@ -139,6 +139,17 @@ fn finish(path: &Path, out: &Path, entry: &str, modules: &[String]) -> ExitCode
         "mypyc".to_string(),
         "--allow-redefinition-new".to_string(),
         "--local-partial-types".to_string(),
+        // Pyfun binds a unit call's result (`_ = f(x)`), which mypy reports
+        // when `f` is annotated `-> None`; the binding is harmless.
+        "--disable-error-code".to_string(),
+        "func-returns-value".to_string(),
+        // At the Python boundary the extern declarations are the contract, and
+        // Pyfun's checker has already held the program to them; mypy's view of
+        // an unannotated module-level value or a stub's union adds nothing.
+        "--disable-error-code".to_string(),
+        "var-annotated".to_string(),
+        "--disable-error-code".to_string(),
+        "union-attr".to_string(),
     ];
     args.extend(modules.iter().map(|m| format!("{m}.py")));
     let status = Command::new(&interpreter)

src/project/mod.rs

@@ -401,7 +401,14 @@ pub fn compile_with(
     if needs_runtime {
         files.push((
             "_pyfun_rt.py".to_string(),
-            python_emitter::emit_for(&lowering::runtime_module(), target),
+            python_emitter::emit_for(
+                &if native {
+                    lowering::runtime_module_native()
+                } else {
+                    lowering::runtime_module()
+                },
+                target,
+            ),
         ));
     }
     Ok(CompiledProject { files, notes })

tests/compile.rs

@@ -4016,6 +4016,32 @@ fn native_annotates_top_level_functions_it_can_spell() {
     assert_eq!(out, ["4.0", "4", "['A']", "Some(1)", "2"]);
 }
 
+#[test]
+fn native_annotates_block_local_functions_and_leaves_unit_returns_bare() {
+    let src = "type Shape = Circle float | Square float\n\
+               let total shapes =\n  \
+                 let area s =\n    \
+                   match s:\n      \
+                     case Circle r: 3.0 * r * r\n      \
+                     case Square w: w * w\n  \
+                 let add acc s = acc + area s\n  \
+                 List.fold add 0.0 shapes\n\
+               let report n = print (n + 0.0)\n\
+               report (total [Circle 1.0, Square 2.0])";
+    let Some((py, out)) = run_native(src) else {
+        return;
+    };
+    assert!(py.contains("    def area(s: Shape) -> float:"), "{py}");
+    assert!(
+        py.contains("    def add(acc: float, s: Shape) -> float:"),
+        "{py}"
+    );
+    // A unit result stays unannotated: `return print(n)` under `-> None` is an
+    // error to mypy.
+    assert!(py.contains("def report(n: float):"), "{py}");
+    assert_eq!(out, ["7.0"]);
+}
+
 #[test]
 fn native_ladders_reach_async_return_bang_matches() {
     // `return! (match …)` in an `async { }` block lowers through its own path,