11from typing import Dict
22
33from mypyc .ir .func_ir import FuncIR
4- from mypyc .ir .ops import Assign , AssignMulti , Box , Branch , Call , CallC , ComparisonOp , DecRef , IncRef , KeepAlive , Return , SetMem , Unbox , Value , Op
4+ from mypyc .ir .ops import (
5+ Assign ,
6+ AssignMulti ,
7+ Box ,
8+ Call ,
9+ CallC ,
10+ ComparisonOp ,
11+ DecRef ,
12+ IncRef ,
13+ KeepAlive ,
14+ Op ,
15+ Return ,
16+ Unbox ,
17+ Value ,
18+ )
519from mypyc .irbuild .ll_builder import LowLevelIRBuilder
620from mypyc .options import CompilerOptions
721from mypyc .transform .ir_transform import IRTransform
@@ -14,6 +28,7 @@ def do_box_unbox_elimination(fn: FuncIR, options: CompilerOptions) -> None:
1428 transform .transform_blocks (fn .blocks )
1529 fn .blocks = builder .blocks
1630
31+
1732def build_use_map (fn : FuncIR ) -> Dict [Value , list [Op ]]:
1833 # Map each Value to a list of ops that use it
1934 use_map : Dict [Value , list [Op ]] = {}
@@ -29,6 +44,7 @@ def build_use_map(fn: FuncIR) -> Dict[Value, list[Op]]:
2944
3045_supported_ops = (Assign , AssignMulti , Call , CallC , ComparisonOp , Return )
3146
47+
3248class BoxUnboxEliminationTransform (IRTransform ):
3349 def __init__ (self , builder : LowLevelIRBuilder , use_map : Dict [Value , list [Op ]]):
3450 super ().__init__ (builder )
@@ -39,7 +55,11 @@ def visit_box(self, op: Box) -> Value | None:
3955 if len (users ) == 0 :
4056 return None
4157 # Check for Unbox->Box
42- if isinstance (op .src , Unbox ) and op .type == op .src .src .type and all (isinstance (user , _supported_ops ) for user in users ):
58+ if (
59+ isinstance (op .src , Unbox )
60+ and op .type == op .src .src .type
61+ and all (isinstance (user , _supported_ops ) for user in users )
62+ ):
4363 unbox = op .src
4464 for user in users :
4565 user .set_sources ([unbox .src if src is op else src for src in user .sources ()])
@@ -51,7 +71,11 @@ def visit_unbox(self, op: Unbox) -> Value | None:
5171 if len (users ) == 0 :
5272 return None
5373 # Check for Box->Unbox
54- if isinstance (op .src , Box ) and op .type == op .src .src .type and all (isinstance (user , _supported_ops ) for user in users ):
74+ if (
75+ isinstance (op .src , Box )
76+ and op .type == op .src .src .type
77+ and all (isinstance (user , _supported_ops ) for user in users )
78+ ):
5579 box = op .src
5680 for user in users :
5781 user .set_sources ([box .src if src is op else src for src in user .sources ()])
0 commit comments