@@ -272,6 +272,12 @@ def _exprs_of_stmt(s: ir.Stmt):
272272 elif isinstance (s , ir .FieldSet ):
273273 yield s .value
274274 yield s .obj
275+ elif isinstance (s , ir .StoreItem ):
276+ yield s .value
277+ yield s .obj
278+ yield s .index
279+ elif isinstance (s , ir .ForObj ):
280+ yield s .iter
275281 elif isinstance (s , ir .ExprStmt ):
276282 yield s .value
277283 elif isinstance (s , ir .If ):
@@ -288,7 +294,7 @@ def _exprs_of_stmt(s: ir.Stmt):
288294def _children (s : ir .Stmt ):
289295 if isinstance (s , ir .If ):
290296 return s .then + s .orelse
291- if isinstance (s , (ir .While , ir .ForRange )):
297+ if isinstance (s , (ir .While , ir .ForRange , ir . ForObj )):
292298 return s .body
293299 return ()
294300
@@ -376,7 +382,8 @@ def _watch_event_stmt() -> str:
376382
377383
378384def _is_effect (e : ir .Expr ) -> bool :
379- if isinstance (e , (ir .GetAttr , ir .CallObject , ir .Truth , ir .CompareObj , ir .ObjToFloat )):
385+ if isinstance (e , (ir .GetAttr , ir .CallObject , ir .Truth , ir .CompareObj , ir .ObjToFloat ,
386+ ir .GetItem , ir .Contains , ir .Format )):
380387 return True
381388 return isinstance (e , ir .BinOp ) and e .type is Type .OBJ
382389
@@ -387,7 +394,7 @@ def _is_effect(e: ir.Expr) -> bool:
387394# Nodes that can run Python code (and so create Python frames the compiled frame count must see):
388395# their slow paths call getattr, the class, __bool__, __float__, a Python callable, a descriptor.
389396_RUNS_PYTHON = (ir .CallObject , ir .GetAttr , ir .Truth , ir .CompareObj , ir .ObjToFloat , ir .FieldGet ,
390- ir .FieldSet , ir .New )
397+ ir .FieldSet , ir .New , ir . GetItem , ir . Contains , ir . Format , ir . StoreItem , ir . ForObj )
391398
392399
393400def _runs_no_python (f : ir .Function ) -> bool :
@@ -403,7 +410,7 @@ def _runs_no_python(f: ir.Function) -> bool:
403410 if any (t is Type .OBJ for t in f .locals .values ()):
404411 return False
405412 for s in _walk_stmts (f .body ):
406- if isinstance (s , ir .FieldSet ):
413+ if isinstance (s , ( ir .FieldSet , ir . StoreItem , ir . ForObj ) ):
407414 return False
408415 for e in _all_exprs (f ):
409416 if isinstance (e , _RUNS_PYTHON ):
@@ -664,7 +671,7 @@ def _check_function(f: ir.Function, functions: dict[str, ir.Function]) -> None:
664671 # a store into a local array is not an effect (ir.py "Local arrays"); into a parameter it is
665672 stores = [s for s in _walk_stmts (f .body )
666673 if (isinstance (s , ir .StoreIndex ) and f .locals .get (s .array ) not in ir .ARRAYS )
667- or isinstance (s , ir .FieldSet )]
674+ or isinstance (s , ( ir .FieldSet , ir . StoreItem , ir . ForObj ) )]
668675 if f .pure and (effects or stores ):
669676 what = type (effects [0 ]).__name__ if effects else type (stores [0 ]).__name__
670677 raise CGenError (f"{ f .name } { what } )
@@ -725,6 +732,8 @@ def _is_closed(f: ir.Function, functions: dict[str, ir.Function], seen=None) ->
725732 if f .name in seen :
726733 return True
727734 seen .add (f .name )
735+ if any (isinstance (s , (ir .StoreItem , ir .ForObj )) for s in _walk_stmts (f .body )):
736+ return False
728737 for e in _all_exprs (f ):
729738 if _is_effect (e ):
730739 return False
@@ -893,6 +902,9 @@ def ev_Const(self, e: ir.Const) -> str:
893902 return self ._own_copy ("Py_None" )
894903 if isinstance (e .value , bool ):
895904 return self ._own_copy ("Py_True" if e .value else "Py_False" )
905+ if type (e .value ) is str and "\0 " not in e .value :
906+ # a str constant: the interned object the module state holds (names table)
907+ return self ._own_copy (f"tp_st->names[{ self .mg .name_index (e .value )} )
896908 raise CGenError (f"{ self .f .name } { e .value !r} )
897909 raise CGenError (f"{ self .f .name } { e .type } )
898910
@@ -934,7 +946,8 @@ def ev_BinOp(self, e: ir.BinOp) -> str:
934946 return t
935947 if lt is Type .OBJ :
936948 t = self .tmp ("PyObject *" )
937- self .emit (f"{ t } { a } { b } { _OBJ_OP [e .op ]} )
949+ op = _OBJ_OP [e .op ] if not e .inplace else f"TP_BINOP_I{ e .op .name }
950+ self .emit (f"{ t } { a } { b } { op } )
938951 self ._own_release (a , b )
939952 self .emit (f"if ({ t } )
940953 return t
@@ -1411,6 +1424,54 @@ def ev_Tuple(self, e: ir.Tuple) -> str:
14111424 self .emit (f"if ({ t } )
14121425 return t
14131426
1427+ def ev_GetItem (self , e : ir .GetItem ) -> str :
1428+ o = self ._obj_operand (e .obj , "GetItem" ) # the container, then the index
1429+ k = self ._obj_operand (e .index , "GetItem index" )
1430+ t = self .tmp ("PyObject *" )
1431+ self .emit (f"{ t } { o } { k } )
1432+ self ._own_release (o , k )
1433+ self .emit (f"if ({ t } )
1434+ return t
1435+
1436+ def ev_Contains (self , e : ir .Contains ) -> str :
1437+ a = self ._obj_operand (e .left , "Contains" ) # CONTAINS_OP: left, then right
1438+ b = self ._obj_operand (e .right , "Contains" )
1439+ t = self .tmp ("int" )
1440+ self .emit (f"{ t } { b } { a } )
1441+ self ._own_release (a , b )
1442+ self .emit (f"if ({ t } )
1443+ if e .negate :
1444+ self .emit (f"{ t } { t } )
1445+ return t
1446+
1447+ def ev_NewList (self , e : ir .NewList ) -> str :
1448+ return self ._own_new ("tp_new_list()" )
1449+
1450+ def ev_NewDict (self , e : ir .NewDict ) -> str :
1451+ return self ._own_new ("tp_new_dict()" )
1452+
1453+ def ev_Format (self , e : ir .Format ) -> str :
1454+ if e .conversion not in - 1 , 97 , 114 , 115 ) or (e .spec is not None and "\0 " in e .spec ):
1455+ raise CGenError (f"{ self .f .name } { e .conversion !r} { e .spec !r} )
1456+ o = self ._obj_operand (e .operand , "Format" )
1457+ spec = "NULL" if e .spec is None else f"tp_st->names[{ self .mg .name_index (e .spec )}
1458+ t = self .tmp ("PyObject *" )
1459+ self .emit (f"{ t } { o } { spec } { e .conversion } )
1460+ self ._own_release (o )
1461+ self .emit (f"if ({ t } )
1462+ return t
1463+
1464+ def ev_BuildString (self , e : ir .BuildString ) -> str :
1465+ if len (e .elements ) < 2 :
1466+ raise CGenError (f"{ self .f .name } { len (e .elements )} )
1467+ items = [self ._obj_operand (x , "BuildString element" ) for x in e .elements ]
1468+ t = self .tmp ("PyObject *" )
1469+ self .emit (f"{{ PyObject *tp_av[{ len (items )} { ', ' .join (items )} )
1470+ self .emit (f" { t } { len (items )} )
1471+ self ._own_release (* items )
1472+ self .emit (f"if ({ t } )
1473+ return t
1474+
14141475 def _assign_array (self , s : ir .Assign , at : Type ) -> None :
14151476 """`local = NewArray | CopyArray`: build into a zeroed temporary, then (only on success)
14161477 free the old array and move the temporary in."""
@@ -1489,6 +1550,48 @@ def st_StoreIndex(self, s: ir.StoreIndex) -> None:
14891550 if not local : # a local array has no list to write back to
14901551 self .emit (f"{ ptr } { slot } )
14911552
1553+ def st_StoreItem (self , s : ir .StoreItem ) -> None :
1554+ v = self ._obj_operand (s .value , "StoreItem value" ) # STORE_SUBSCR: value, container, index
1555+ o = self ._obj_operand (s .obj , "StoreItem" )
1556+ k = self ._obj_operand (s .index , "StoreItem index" )
1557+ self .emit (f"tp_s = tp_setitem({ o } { k } { v } )
1558+ self ._own_release (o , k , v )
1559+ self .emit ("if (tp_s < 0) { tp_rc = -1; goto tp_exit; }" )
1560+
1561+ def st_ForObj (self , s : ir .ForObj ) -> None :
1562+ if not s .targets or len (set (s .targets )) != len (s .targets ) or (not s .unpack and len (s .targets ) != 1 ):
1563+ raise CGenError (f"{ self .f .name } { s .targets !r} { s .unpack } )
1564+ for name in s .targets :
1565+ if self .types .get (name ) is not Type .OBJ :
1566+ raise CGenError (f"{ self .f .name } { name } )
1567+ o = self ._obj_operand (s .iter , "ForObj iterable" )
1568+ it = self .tmp ("PyObject *" ) # the iterator: owned by this statement across iterations
1569+ self .emit (f"{ it } { o } ) # GET_ITER
1570+ self ._own_release (o )
1571+ self .emit (f"if ({ it } )
1572+ item = self .tmp ("PyObject *" )
1573+ self .emit ("for (;;) {" )
1574+ self .depth += 1
1575+ self .poll ()
1576+ self .emit (f"{ item } { it } ) # FOR_ITER (StopIteration ends the loop)
1577+ self .emit (f"if ({ item } )
1578+ if not s .unpack :
1579+ self ._own_move (_ident ("l" , s .targets [0 ]), item )
1580+ else :
1581+ n = len (s .targets )
1582+ self .emit (f"{{ PyObject *tp_u[{ n } )
1583+ self .emit (f" tp_s = tp_unpack({ item } { n } ) # UNPACK_SEQUENCE; nothing owned on error
1584+ self ._own_release (item )
1585+ self .emit (" if (tp_s < 0) { tp_rc = -1; goto tp_exit; }" )
1586+ for i , name in enumerate (s .targets ):
1587+ local = _ident ("l" , name )
1588+ self .emit (f" {{ PyObject *tp_old = { local } { local } { i } )
1589+ self .emit ("}" )
1590+ self .stmts (s .body )
1591+ self .depth -= 1
1592+ self .emit ("}" )
1593+ self ._own_release (it ) # after the loop and after a `break`; a return/error exits via tp_exit
1594+
14921595 def st_FieldSet (self , s : ir .FieldSet ) -> None :
14931596 slot = self .mg .slot (self .f .name , s .cls , s .field )
14941597 v = self ._obj_operand (s .value , "FieldSet value" ) # CPython: value, then the object
0 commit comments