[Enhancement] Resolve reference cycle #795
Conversation
WalkthroughAdds nonlocal for outer_stack inside a decorator wrapper, clears outer_stack when functions are defined in a class and also unconditionally before parsing non-class functions, preserving TypeError for non-function inputs. Ensures outer_stack can be rebound and avoids leakage of stack state across decorated functions. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor Dev as Developer
participant Dec as decorator_wrapper
participant Util as utils.is_defined_in_class
participant Func as func
note over Dec: nonlocal outer_stack
Dev->>Dec: Call decorator on func
Dec->>Dec: Validate input (TypeError if not function)
Dec->>Util: Check if func defined in class
alt Defined in class
Dec->>Dec: outer_stack = None
Dec-->>Dev: Return original func
else Not defined in class
Dec->>Dec: outer_stack = None
Dec->>Dec: Parse/handle func with cleared stack
Dec-->>Dev: Return wrapped func
end
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Pre-merge checks (1 warning, 2 inconclusive)❌ Failed checks (1 warning, 2 inconclusive)
Poem
Tip 👮 Agentic pre-merge checks are now available in preview!Pro plan users can now enable pre-merge checks in their settings to enforce checklists before merging PRs.
Example: reviews:
pre_merge_checks:
custom_checks:
- name: "Undocumented Breaking Changes"
mode: "warning"
instructions: |
Flag potential breaking changes that are not documented:
1. Identify changes to public APIs/exports, CLI flags, environment variables, configuration keys, database schemas, or HTTP/GraphQL endpoints (including removed/renamed items and changes to types, required params, return values, defaults, or behavior).
2. Ignore purely internal/private changes (e.g., code not exported from package entry points or marked internal).
3. Verify documentation exists: a "Breaking Change" section in the PR description and updates to CHANGELOG.md.Please share your feedback with us on this Discord post. ✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @LeiWang1999, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request addresses a potential memory management issue by resolving a reference cycle. This is achieved by explicitly clearing a non-local variable (outer_stack) within a decorator, ensuring that references are properly released and preventing memory leaks.
Highlights
- Reference Cycle Resolution: Explicitly sets the
outer_stackvariable toNonewithin theprim_funcdecorator'sdecorator_wrapperto break potential reference cycles and aid garbage collection.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
There was a problem hiding this comment.
Code Review
This pull request addresses a potential reference cycle in the prim_func decorator by ensuring the outer_stack variable is cleared after use. The approach is functionally correct. My review includes a suggestion to refactor the implementation slightly to improve code clarity and remove redundancy, making the logic for clearing the stack reference more straightforward.
| if utils.is_defined_in_class(outer_stack, func): | ||
| outer_stack = None | ||
| return func | ||
| outer_stack = None |
There was a problem hiding this comment.
While this correctly clears outer_stack to prevent reference cycles, the logic can be simplified to avoid setting outer_stack = None in two separate code paths. You can determine whether the function is defined in a class, clear outer_stack immediately after its use, and then return if necessary. This makes the code more concise and easier to follow.
| if utils.is_defined_in_class(outer_stack, func): | |
| outer_stack = None | |
| return func | |
| outer_stack = None | |
| is_in_class = utils.is_defined_in_class(outer_stack, func) | |
| outer_stack = None | |
| if is_in_class: | |
| return func |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tilelang/language/tir/entry.py (1)
36-39: Micro: consider lighter stack capture if compatible.
inspect.stack()is expensive and retains source context. Ifutils.is_defined_in_classdoesn’t need context lines,inspect.stack(context=0)can reduce overhead and memory. Please verifyutils.is_defined_in_class’s requirements before changing.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
tilelang/language/tir/entry.py(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: bot-task
🔇 Additional comments (1)
tilelang/language/tir/entry.py (1)
43-43: Good fix: nonlocal prevents UnboundLocalError and enables cleanup.Declaring
nonlocal outer_stackis necessary since you rebindouter_stackbelow. This prevents a local-scope shadowing bug and lets you clear the captured frames.
| if not inspect.isfunction(func): | ||
| raise TypeError(f"Expect a function, but got: {func}") | ||
| nonlocal outer_stack | ||
| if utils.is_defined_in_class(outer_stack, func): | ||
| outer_stack = None | ||
| return func | ||
| outer_stack = None | ||
| f = parse(func, utils.inspect_function_capture(func), check_well_formed=check_well_formed) |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Consolidate cleanup with try/finally; also clear on TypeError path.
Two explicit resets plus an early raise leave one path (TypeError) where the captured frames aren’t cleared. Use a single finally to guarantee cleanup across all branches and exceptions, and move nonlocal to the top.
Apply:
def decorator_wrapper(func):
- if not inspect.isfunction(func):
- raise TypeError(f"Expect a function, but got: {func}")
- nonlocal outer_stack
- if utils.is_defined_in_class(outer_stack, func):
- outer_stack = None
- return func
- outer_stack = None
- f = parse(func, utils.inspect_function_capture(func), check_well_formed=check_well_formed)
- setattr(f, "__name__", func.__name__) # noqa: B010
- return f
+ nonlocal outer_stack
+ if not inspect.isfunction(func):
+ outer_stack = None
+ raise TypeError(f"Expect a function, but got: {func}")
+ try:
+ if utils.is_defined_in_class(outer_stack, func):
+ return func
+ f = parse(func, utils.inspect_function_capture(func), check_well_formed=check_well_formed)
+ setattr(f, "__name__", func.__name__) # noqa: B010
+ return f
+ finally:
+ outer_stack = None📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if not inspect.isfunction(func): | |
| raise TypeError(f"Expect a function, but got: {func}") | |
| nonlocal outer_stack | |
| if utils.is_defined_in_class(outer_stack, func): | |
| outer_stack = None | |
| return func | |
| outer_stack = None | |
| f = parse(func, utils.inspect_function_capture(func), check_well_formed=check_well_formed) | |
| def decorator_wrapper(func): | |
| nonlocal outer_stack | |
| if not inspect.isfunction(func): | |
| outer_stack = None | |
| raise TypeError(f"Expect a function, but got: {func}") | |
| try: | |
| if utils.is_defined_in_class(outer_stack, func): | |
| return func | |
| f = parse( | |
| func, | |
| utils.inspect_function_capture(func), | |
| check_well_formed=check_well_formed, | |
| ) | |
| setattr(f, "__name__", func.__name__) # noqa: B010 | |
| return f | |
| finally: | |
| outer_stack = None |
🧰 Tools
🪛 Ruff (0.12.2)
42-42: Avoid specifying long messages outside the exception class
(TRY003)
🤖 Prompt for AI Agents
In tilelang/language/tir/entry.py around lines 41 to 48, the code currently
declares nonlocal and resets outer_stack in multiple places but raises TypeError
before clearing captured frames; move the nonlocal outer_stack declaration to
the top of the function block and refactor the logic to use a try/finally so
outer_stack is always reset to None in the finally clause (remove the duplicate
explicit resets and ensure the TypeError path also runs the finally cleanup);
maintain existing behavior of returning early when
utils.is_defined_in_class(...) is true and preserve raising the TypeError when
func is not a function.
Co-authored-by: Huanqi Cao <caohuanqi@deepseek.com>
1 participant
Summary by CodeRabbit
New Features
Bug Fixes
Refactor