feat(arxjit): add Python-subset validation for decorated functions

[Jaskirat-s7]

Summary

Sprint 2, Python-subset validation (wiki "Proposed Issue 4: Define and validate the supported pure-Python subset"). Builds on source extraction (#96/#98): it adds arxjit.validation, the stage that runs after extract_source, walking the parsed function and accepting only the v1 supported subset while rejecting everything else with clear, located diagnostics.

Issue 4 checklist coverage:

• Implement a validation layer for Python AST nodes
• Accept only supported built-in Python constructs
• Reject unsupported syntax with clear diagnostics

What it accepts (the v1 subset)

Function definitions, typed scalar arguments, int/float/bool literals, single-target local assignments, arithmetic (+ - * / // % **), comparisons (including chained), boolean and/or, unary + - not, if/else, while, for ... in range(...), return, pass, and docstrings.

Everything else is rejected: classes, async, closures and globals, lambdas, containers and comprehensions, calls (except range(...) in a for), imports, exceptions, with, attribute access, subscripting, f-strings, walrus, break/continue, loop else, generic (PEP 695) functions, and *args/**kwargs/keyword-only/default parameters.

How it works

validate(extracted: ExtractedSource) -> None raises UnsupportedSyntaxError carrying one Diagnostic per rejected construct, so a function with several problems reports all of them at once instead of one-per-run.

Type-based dispatch. The validator dispatches on node type with plum @dispatch (a single visit method overloaded per node type), matching the visitor convention used elsewhere in the Arx packages. The ast.AST overload is the fail-closed default: any node without an accepting overload is rejected, either with a specific message from a lookup table or a generic fallback, so unhandled or future syntax fails safe rather than passing silently. A rejected node's children are not descended into (one diagnostic per bad subtree), while independent sibling statements are each validated in full.

Binding analysis. validate collects the names bound in the function's own scope — parameters plus every assignment and loop target — and rejects any name read that is not bound, which is how closure captures and module globals are caught (an extracted nested function keeps its free variables as plain Name loads, so this cannot be done structurally). The collector deliberately does not descend into nested functions, lambdas, or classes: those have their own locals, and collecting them would both mask free-variable reads in this scope and misreport a nested binding as shadowing an outer name. A nested definition's own name does bind in the enclosing scope, so it is recorded while its body is skipped. The tuple, list, starred, annotated, augmented, and walrus binding forms are all recorded, so a statement the subset rejects does not additionally report a spurious free variable for the name it binds.

range must really be range. A for loop's iterable must be a call to the builtin range with 1–3 positional arguments and no keywords. Because only one of the three ways the name can be shadowed is visible in the AST, extraction preserves the defining module's namespace and the function's captured names, and validation rejects the loop when range resolves to:

shadowed by	detected via
a local variable or argument	the binding analysis above
an enclosing function's variable (a closure capture)	the captured names (co_freevars)
a module-level name	the defining module's namespace vs builtins.range

Each of those is a TypeError in Python when the shadowing value is not callable, so compiling them as the range intrinsic would silently disagree with the Python fallback.

Locations. Diagnostic columns are converted from ast's zero-based UTF-8 byte offsets to the one-based character contract from #95, and node line numbers are resolved against the real file — correctly for nested functions, where node.lineno is a real file line but the extracted source is indexed from its own first line.

Scope of this PR

To keep review small, this PR is the validator plus a test set that exercises every code branch (100% coverage): the accept path, collect-all-at-once, real-file line/column correctness including a nested-function case, the binding and range rules, argument-shape rules, and the helpers. The exhaustive one-test-per-unsupported-construct rejection matrix lands in a small stacked follow-up — it is completeness coverage over the whitelist, and all of it routes through code paths already covered here.

Design points to review

1. Diagnostic code stays None. I'd like to introduce a stable code scheme (e.g. AJ001) with this validation work, since it is the first stage emitting many distinct diagnostics — flagging for a decision on format/prefix before baking numbers in.
2. Standalone expression statements are rejected (a value computed and discarded), with docstrings and pass explicitly allowed; a bare yield is routed to the specific "generators are not supported" message.
3. Loop else is rejected, consistent with rejecting break: a loop else only runs when the loop exits without break, so accepting it would be dead code and would complicate lowering.
4. validate takes ExtractedSource, returns None, raises on failure — the fallback-vs-strict behavior belongs to the @jit wiring PR, not here.
5. When the module namespace is unavailable — only possible for a hand-built ExtractedSource, since extract_source always provides it for a real function — module-level shadowing cannot be observed and the name is taken to be the builtin. Happy to make that path stricter if you would rather it fail closed.

Verification

• 39 validation tests (103 in arxjit), coverage 100%. Location assertions read this test file itself, so they don't rot when it is edited.
• The ast semantics relied on (byte-vs-character columns, FunctionDef.lineno, version-gated node types) were probed on CPython 3.10, 3.11, 3.13 and 3.14, and the binding, range-shadowing and generic-function rules were verified to behave identically on all four. TryStar (except*, 3.11+) and type_params (PEP 695, 3.12+) are both accessed through runtime guards, so the module imports on 3.10.
• All gates green locally: pytest + the 86% coverage gate, ruff format/check, mypy strict, bandit, vulture, mccabe, and douki sync idempotent.

Note on the diff size

GitHub reports a large addition count because both files are new. Roughly half of it is mandatory Douki docstrings plus blank lines; the validator logic itself is about 250 lines, alongside a flat construct-to-message lookup table that reads as data.

[xmnlab]

thanks for working on that @Jaskirat-s7

for function names, please use everything in lowercase. for visit methods use multiple dispatch from plumb instead of creating visit_*

other than that, I found three cases where unsupported semantics can pass the new validator.

1. High — closures and global-variable reads are accepted

visit_Name() accepts every variable reference without checking whether the name is an argument or a local variable. Validation then walks only the body after checking argument shape; it performs no scope or binding analysis.

This closure therefore validates successfully:

def outer(scale):
def kernel(x: int) -> int:
return x * scale

return kernel

validate(extract_source(outer(2)))

The extracted node contains only kernel, so there is no nested FunctionDef or nonlocal node to reject. scale is simply an accepted ast.Name.

Module globals have the same problem:

SCALE = 2

def kernel(x: int) -> int:
return x * SCALE

This contradicts the stated rejection of closures and the pure-function model. It could also leave lowering with an undefined value or cause compiled behavior to differ from the Python fallback.

The validator needs name-binding analysis, or ExtractedSource/validate() must retain enough function metadata to detect co_freevars and globals.

Location: packages/arxjit/src/arxjit/validation.py, visit_Name() and validate().

2. High — a shadowed or invalid range call is treated as the builtin

_is_range_call() checks only that the callee is a bare name spelled range and that there are no keyword arguments. It does not detect shadowing or validate the number of positional arguments. visit_For() then accepts the call and validates only its arguments.

These all pass:

def kernel(range: int) -> int:
total = 0
for i in range(3):
total = total + i
return total

def kernel(n: int) -> int:
range = n
for i in range(3):
pass
return 0

def kernel() -> int:
for i in range():
pass
return 0

def kernel() -> int:
for i in range(1, 2, 3, 4):
pass
return 0

The first two call an integer in Python; the latter two raise TypeError. A compiler that treats the syntax as an intrinsic could silently change behavior.

At minimum, reject local or argument bindings named range and require one to three positional arguments. Module-level shadowing requires either access to the original function’s globals or an explicit language rule that reserves range.

Location: packages/arxjit/src/arxjit/validation.py, _is_range_call() and visit_For().

3. Medium — Python 3.12+ generic functions bypass the fail-closed policy

The PR says new Python syntax fails closed, but the top-level FunctionDef itself is never visited. validate() checks async status, argument shapes, and body statements only. Argument and return annotations are also deliberately not inspected.

On Python 3.12+, this can therefore validate:

def identity[T](x: T) -> T:
return x

Its type_params field is never checked, and all body nodes are accepted. This is a concrete hole in the claimed fail-closed handling of future/version-specific syntax.

A guarded check such as if getattr(node, "type_params", ()): should reject generic function parameters. Annotation reconciliation can remain a later stage if that is intentional.

Location: packages/arxjit/src/arxjit/validation.py, validate().

Summary

Changes requested. The AST whitelist is well structured and the diagnostic-location handling is thoughtful, but the missing binding analysis means closures, globals, and shadowed range calls currently cross the validation boundary despite being outside the declared subset.

[yuvimittal]

@Jaskirat-s7 , It looks like diagnostic.column is still None where the test expects a non-None value on 3.13+.

Could you take a look and confirm whether this is related to the changes in this PR or a pre-existing issue?

[Jaskirat-s7]

@yuvimittal Thanks for flagging it , I dug into the CI logs and it's a pre-existing issue, not from this PR. The failing test is test_malformed_fstring_column_is_validated in test_source.py (merged in #98), not a validation test — all the validation tests pass on 3.13/3.14.

what i noticed was that the newer 3.13 CI runner produces a different error for a malformed f-string ("invalid syntax. Perhaps you forgot a comma?"), and its offset points into a synthetic replacement-field string rather than the real source line. Our _column_of guard correctly detects that the offset isn't trustworthy and returns column=None — that's the intended behavior (better no column than a misleading one). So the production code is right.

The problem is the test assumption: it asserts column is not None for 3.13+, which was a bet on CPython patch behavior. When #98 merged, the 3.13 runner was an older patch that gave a validated column; it's since updated. I'll relax that assertion to accept either a validated column or None on 3.13+ (the same tolerant check I already use for 3.12), so it stops depending on the exact CPython patch. I'll fold that fix into this PR since it's what's reddening CI here.

[Jaskirat-s7]

Thanks for the thorough review @xmnlab all four addressed in the latest commit:

1. Dispatch + naming. Converted the visitor to plum @dispatch on a single lowercase visit method, matching the convention in irx/base/visitors/base.py (one visit, overloaded by node type, fail-closed visit(ast.AST) default). That also removes the CapWords visit_* names. Added plum-dispatch as a direct arxjit dependency.

2. Closures + globals (High) , to this i added binding analysis: validate collects the bound names (arguments + every assignment/for-loop target), and a Name read that isn't bound is rejected as a free variable. Your exact outer(scale)/kernel closure and the SCALE global example now both fail with "reference to 'scale', which is not a local variable or argument (closures and globals are not supported)".

3. Shadowed/invalid range (High). range in a for is now rejected when it's shadowed by a local/argument, when keyword arguments are used, or when it doesn't have 1–3 positional arguments. All four of your snippets are handled.

4. PEP 695 generics (Medium), for this i added a guarded type_params check, so def identityT is rejected on 3.12+ (and the guard keeps 3.10 importing).

i have verified the closed holes and the accepted subset on CPython 3.10, 3.11, 3.13, and 3.14. Also dropped an unreachable BoolOp operator check (Python only has and/or). Coverage stays 100%.

[xmnlab]

thanks @Jaskirat-s7 for addressing the points I made before.

I found two remaining issues.

1. High — a module-global range shadow still bypasses validation

in the comments it is mentioned tge shadowed range is rejected when it is a local variable or argument. That is accurate, but the earlier review also separately noted module-level shadowing.

The implementation only checks:

if "range" in self.bound:
...

self.bound contains parameters and local assignment or loop targets. It contains no information about the original function’s globals. Also, _visit_for_iter() does not visit the range callee as an ordinary Name, so the general free-variable check never sees it.

This therefore still validates:

range = 42

def kernel(n: int) -> int:
total = 0
for i in range(n):
total = total + i
return total

Python raises TypeError because the module global shadows the builtin. A compiled implementation treating it as the range intrinsic would silently change the function’s behavior.

The current tests cover argument shadowing, arity, and keyword arguments, but not module-global shadowing.

Possible resolutions:

Preserve the original callable or its global namespace through extraction and verify that fn.globals.get("range", builtins.range) is builtins.range.

Move this check to the @jit wiring stage, where the callable is still available.

Explicitly define range as a reserved intrinsic in ArxJIT, though that creates a deliberate difference from Python fallback semantics and must be documented and enforced consistently.

2. Medium — binding analysis incorrectly descends into nested scopes

_bound_names() uses ast.walk(node), which traverses nested functions, lambdas, and classes. Assignment and loop targets found inside those nested scopes are therefore incorrectly treated as locals of the outer function.

For example:

def kernel(n: int) -> int:
def helper() -> int:
range = 10
return range

total = 0
for i in range(n):
    total = total + i
return total

The outer function’s range is not locally shadowed, but _bound_names() collects range from helper() and produces the misleading diagnostic:

range is shadowed by a local variable or argument

A similar case can hide an independent free-variable diagnostic:

def kernel() -> int:
def helper() -> int:
scale = 2
return scale

return scale

The nested function is correctly rejected, but the outer scale read is incorrectly considered bound, undermining the PR’s collect-all-at-once guarantee.

The binding collector should traverse only the current function’s scope and skip the bodies of nested FunctionDef, AsyncFunctionDef, Lambda, and ClassDef nodes.

Non-blocking documentation

The PR description still describes the original ast.NodeVisitor implementation and says there are 25 tests. It should be updated to describe Plum dispatch, binding analysis, the revised range rules, the generic-function check, and the current test count.

summary:

Changes requested. The original review findings are substantially and thoughtfully addressed, and CI is green. Module-global range shadowing remains a behavior-changing validation bypass, while nested-scope traversal makes the new binding analysis produce inaccurate or incomplete diagnostics.

[Jaskirat-s7]

Thanks @xmnlab , both addressed, plus a third case your findings led me to.

1. Module-global range shadow (high). Took your first option: extraction now preserves the defining module's namespace on ExtractedSource, and validation verifies range resolves to builtins.range. Your range = 42 example is rejected; a clean namespace, or range explicitly bound to the builtin, still validates.

2. Nested-scope binding leak (medium). The collector now traverses only the current scope, skipping the bodies of FunctionDef, AsyncFunctionDef, Lambda and ClassDef while still recording a nested definition's own name (it does bind in the enclosing scope). Both of your examples are fixed: the helper-binds-range case no longer claims "range is shadowed", and the masked-scale case now reports both problems instead of one. I also record the tuple/list/starred/annotated/augmented/walrus binding forms, so a rejected statement doesn't additionally emit a spurious free-variable diagnostic for the name it binds.

3. Stale description. Updated , plum dispatch, binding analysis, the revised range rules, the generic-function check, and the current test counts.

One more I found while auditing the above: a closure-captured range slipped through both checks , it's neither a local nor a module global, but Python resolves it through the closure cell and raises TypeError:

def outer():
range = 5
def kernel(n):
for i in range(n): ...
It was uniquely exposed because range is exempt from the free-variable check by design. Extraction now also preserves co_freevars, and all three shadowing paths are consolidated into one check that reports which applies — a local/argument, an enclosing function's variable, or a module-level name. Verified on 3.10/3.11/3.13/3.14.

One thing I left deliberately and flagged as design point 5 ,when the module namespace is unavailable (only for a hand-built ExtractedSource; extract_source always provides it) module-level shadowing can't be observed and the name is assumed to be the builtin. Happy to make that fail closed if you'd prefer.

[xmnlab]

thanks @Jaskirat-s7

The latest changes address the previous blockers: module-level and closure-captured "range" shadowing are detected, binding collection respects nested scopes, the PR description is current, regression tests cover the reported cases, and CI is green.

I found one remaining edge case involving functions created with a customised "builtins" mapping, plus a diagnostic-quality issue for some unsupported binding forms. Neither needs to block this PR in my view.

Please track a follow-up to either preserve and validate the function’s actual builtins namespace or explicitly reject non-standard builtins before treating "range" as an intrinsic. A separate follow-up can complete binding collection for unsupported constructs such as imports and "with ... as".

Approved from my side.