3030_TRACER = trace .get_tracer (__name__ )
3131
3232
33+ def _normalize_global_let_name (name : str ) -> str :
34+ return name if name .startswith ("$" ) else f"${ name }
35+
36+
3337def span (frame_index : int = 0 ) -> bindings .RustSpan :
3438 """
3539 Returns a span for the current file and line.
@@ -44,10 +48,6 @@ def span(frame_index: int = 0) -> bindings.RustSpan:
4448 return bindings .RustSpan ("" , 0 , 0 )
4549
4650
47- def _normalize_global_let_name (name : str ) -> str :
48- return name if name .startswith ("$" ) else f"${ name }
49-
50-
5151@dataclass
5252class EGraphState :
5353 """
@@ -107,6 +107,14 @@ def copy(self) -> EGraphState:
107107 def _run_program (self , * commands : bindings ._Command ) -> list [bindings ._CommandOutput ]:
108108 return call_with_current_trace (self .egraph .run_program , * commands )
109109
110+ @staticmethod
111+ def _persistent_scheduler_name (scheduler : BackOffDecl ) -> str :
112+ return f"_persistent_scheduler_{ scheduler .id .hex }
113+
114+ @staticmethod
115+ def _local_scheduler_name (index : int ) -> str :
116+ return f"_scheduler_{ index }
117+
110118 @_TRACER .start_as_current_span ("run_schedule_to_egg" )
111119 def run_schedule_to_egg (self , schedule : ScheduleDecl ) -> bindings ._Command :
112120 """
@@ -115,17 +123,19 @@ def run_schedule_to_egg(self, schedule: ScheduleDecl) -> bindings._Command:
115123 If there exists any custom schedulers in the schedule, it will be turned into a custom extract command otherwise
116124 will be a normal run command.
117125 """
118- processed_schedule = self ._process_schedule (schedule )
126+ processed_schedule , persistent_schedulers = self ._process_schedule (schedule )
119127 if processed_schedule is None :
120128 return bindings .RunSchedule (self ._schedule_to_egg (schedule ))
129+ for scheduler in persistent_schedulers :
130+ self ._run_program (self ._persistent_scheduler_to_egg (scheduler ))
121131 top_level_schedules = self ._schedule_with_scheduler_to_egg (processed_schedule , [])
122132 if len (top_level_schedules ) == 1 :
123133 schedule_expr = top_level_schedules [0 ]
124134 else :
125135 schedule_expr = bindings .Call (span (), "seq" , top_level_schedules )
126136 return bindings .UserDefined (span (), "run-schedule" , [schedule_expr ])
127137
128- def _process_schedule (self , schedule : ScheduleDecl ) -> ScheduleDecl | None :
138+ def _process_schedule (self , schedule : ScheduleDecl ) -> tuple [ ScheduleDecl | None , tuple [ BackOffDecl , ...]] :
129139 """
130140 Processes a schedule to determine if it contains any custom schedulers.
131141
@@ -134,19 +144,23 @@ def _process_schedule(self, schedule: ScheduleDecl) -> ScheduleDecl | None:
134144
135145 Also processes all rulesets in the schedule to make sure they are registered.
136146 """
137- bound_schedulers : list [UUID ] = []
147+ bound_schedulers : list [BackOffDecl ] = []
138148 unbound_schedulers : list [BackOffDecl ] = []
149+ persistent_schedulers : dict [UUID , BackOffDecl ] = {}
139150
140151 def helper (s : ScheduleDecl ) -> None :
141152 match s :
142153 case LetSchedulerDecl (scheduler , inner ):
143- bound_schedulers .append (scheduler . id )
154+ bound_schedulers .append (scheduler )
144155 return helper (inner )
145156 case RunDecl (ruleset_name , _, scheduler ):
146157 self .ruleset_to_egg (ruleset_name )
147- if scheduler and scheduler .id not in bound_schedulers :
148- unbound_schedulers .append (scheduler )
149- case SaturateDecl (inner , _) | RepeatDecl (inner , _):
158+ if scheduler and scheduler .id not in s .id for s in bound_schedulers }:
159+ if scheduler .persistent :
160+ persistent_schedulers [scheduler .id ] = scheduler
161+ else :
162+ unbound_schedulers .append (scheduler )
163+ case SaturateDecl (inner ) | RepeatDecl (inner , _):
150164 return helper (inner )
151165 case SequenceDecl (schedules ):
152166 for sc in schedules :
@@ -156,16 +170,16 @@ def helper(s: ScheduleDecl) -> None:
156170 return None
157171
158172 helper (schedule )
159- if not bound_schedulers and not unbound_schedulers :
160- return None
173+ if not bound_schedulers and not unbound_schedulers and not persistent_schedulers :
174+ return None , ()
161175 for scheduler in unbound_schedulers :
162176 schedule = LetSchedulerDecl (scheduler , schedule )
163- return schedule
177+ return schedule , tuple ( persistent_schedulers . values ())
164178
165179 def _schedule_to_egg (self , schedule : ScheduleDecl ) -> bindings ._Schedule :
166180 msg = "Should never reach this, let schedulers should be handled by custom scheduler"
167181 match schedule :
168- case SaturateDecl (schedule , _ ):
182+ case SaturateDecl (schedule ):
169183 return bindings .Saturate (span (), self ._schedule_to_egg (schedule ))
170184 case RepeatDecl (schedule , times ):
171185 return bindings .Repeat (span (), times , self ._schedule_to_egg (schedule ))
@@ -184,33 +198,40 @@ def _schedule_to_egg(self, schedule: ScheduleDecl) -> bindings._Schedule:
184198 assert_never (schedule )
185199
186200 def _schedule_with_scheduler_to_egg ( # noqa: C901, PLR0912
187- self , schedule : ScheduleDecl , bound_schedulers : list [UUID ]
201+ self , schedule : ScheduleDecl , bound_schedulers : list [BackOffDecl ]
188202 ) -> list [bindings ._Expr ]:
189203 """
190204 Turns a scheduler into an egg expression, to be used with a custom extract command.
191205
192206 The bound_schedulers is a list of all the schedulers that have been bound. We can lookup their name as `_scheduler_{index}`.
193207 """
194208 match schedule :
195- case LetSchedulerDecl (BackOffDecl (id , match_limit , ban_length , egg_like ), inner ):
196- name = f"_scheduler_{ len (bound_schedulers )}
197- bound_schedulers .append (id )
209+ case LetSchedulerDecl (scheduler , inner ):
210+ match_limit = scheduler .match_limit
211+ ban_length = scheduler .ban_length
212+ fresh_rematch = scheduler .fresh_rematch
213+ name = self ._local_scheduler_name (len (bound_schedulers ))
214+ bound_schedulers .append (scheduler )
198215 args : list [bindings ._Expr ] = []
199216 if match_limit is not None :
200217 args .append (bindings .Var (span (), ":match-limit" ))
201218 args .append (bindings .Lit (span (), bindings .Int (match_limit )))
202219 if ban_length is not None :
203220 args .append (bindings .Var (span (), ":ban-length" ))
204221 args .append (bindings .Lit (span (), bindings .Int (ban_length )))
205- scheduler_name = "back-off-egg " if egg_like else "back-off"
222+ scheduler_name = "back-off-fresh " if fresh_rematch else "back-off"
206223 back_off_decl = bindings .Call (span (), scheduler_name , args )
207224 let_decl = bindings .Call (span (), "let-scheduler" , [bindings .Var (span (), name ), back_off_decl ])
208225 return [let_decl , * self ._schedule_with_scheduler_to_egg (inner , bound_schedulers )]
209226 case RunDecl (ruleset_ident , until , scheduler ):
210227 args = [bindings .Var (span (), str (ruleset_ident ))]
211228 if scheduler :
212229 name = "run-with"
213- scheduler_name = f"_scheduler_{ bound_schedulers .index (scheduler .id )}
230+ scheduler_name = self ._persistent_scheduler_name (scheduler )
231+ for i , bound in enumerate (bound_schedulers ):
232+ if bound .id == scheduler .id :
233+ scheduler_name = self ._local_scheduler_name (i )
234+ break
214235 args .insert (0 , bindings .Var (span (), scheduler_name ))
215236 else :
216237 name = "run"
@@ -225,10 +246,8 @@ def _schedule_with_scheduler_to_egg( # noqa: C901, PLR0912
225246 raise ValueError (msg )
226247 args .append (fact_egg .expr )
227248 return [bindings .Call (span (), name , args )]
228- case SaturateDecl (inner , stop_when_no_updates ):
249+ case SaturateDecl (inner ):
229250 args = self ._schedule_with_scheduler_to_egg (inner , bound_schedulers )
230- if stop_when_no_updates :
231- args = [bindings .Var (span (), ":stop-when-no-updates" ), * args ]
232251 return [bindings .Call (span (), "saturate" , args )]
233252 case RepeatDecl (inner , times ):
234253 return [
@@ -249,6 +268,22 @@ def _schedule_with_scheduler_to_egg( # noqa: C901, PLR0912
249268 case _:
250269 assert_never (schedule )
251270
271+ def _persistent_scheduler_to_egg (self , scheduler : BackOffDecl ) -> bindings ._Command :
272+ args : list [bindings ._Expr ] = []
273+ if scheduler .match_limit is not None :
274+ args .append (bindings .Var (span (), ":match-limit" ))
275+ args .append (bindings .Lit (span (), bindings .Int (scheduler .match_limit )))
276+ if scheduler .ban_length is not None :
277+ args .append (bindings .Var (span (), ":ban-length" ))
278+ args .append (bindings .Lit (span (), bindings .Int (scheduler .ban_length )))
279+ scheduler_name = "back-off-fresh" if scheduler .fresh_rematch else "back-off"
280+ back_off_decl = bindings .Call (span (), scheduler_name , args )
281+ return bindings .UserDefined (
282+ span (),
283+ "let-scheduler" ,
284+ [bindings .Var (span (), self ._persistent_scheduler_name (scheduler )), back_off_decl ],
285+ )
286+
252287 def ruleset_to_egg (self , ident : Ident ) -> None :
253288 """
254289 Registers a ruleset if it's not already registered.
0 commit comments