islpy 2025.1.5__cp311-cp311-macosx_11_0_arm64.whl → 2025.2.1__cp311-cp311-macosx_11_0_arm64.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
islpy/_monkeypatch.py ADDED
@@ -0,0 +1,1097 @@
1
+ import os
2
+ import re
3
+ from collections.abc import Callable, Collection, Iterable, Mapping, Sequence
4
+ from functools import update_wrapper
5
+ from sys import intern
6
+ from typing import (
7
+ TYPE_CHECKING,
8
+ Any,
9
+ ClassVar,
10
+ Concatenate,
11
+ Literal,
12
+ ParamSpec,
13
+ Protocol,
14
+ TypeAlias,
15
+ TypeVar,
16
+ cast,
17
+ )
18
+ from warnings import warn
19
+
20
+
21
+ if TYPE_CHECKING:
22
+ import islpy._isl as _isl
23
+ else:
24
+ import sys
25
+ if "_isl" not in sys.modules:
26
+ import islpy._isl as _isl
27
+ else:
28
+ # This is used for monkeypatching during stub generation.
29
+ # See stubgen/stubgen.py and CMakeLists for orchestration details.
30
+ import _isl
31
+
32
+
33
+ ALL_CLASSES: tuple[type, ...] = tuple(
34
+ getattr(_isl, cls) for cls in dir(_isl) if cls[0].isupper())
35
+ EXPR_CLASSES: tuple[type, ...] = tuple(cls for cls in ALL_CLASSES
36
+ if "Aff" in cls.__name__ or "Polynomial" in cls.__name__)
37
+ ARITH_CLASSES: tuple[type, ...] = (
38
+ _isl.Aff, _isl.PwAff, _isl.QPolynomial, _isl.PwQPolynomial)
39
+
40
+ _CHECK_DIM_TYPES: tuple[_isl.dim_type, ...] = (
41
+ _isl.dim_type.in_, _isl.dim_type.param, _isl.dim_type.set)
42
+
43
+
44
+ # {{{ typing helpers
45
+
46
+ SelfT = TypeVar("SelfT")
47
+
48
+ BasicT = TypeVar("BasicT", _isl.BasicSet, _isl.BasicMap)
49
+
50
+ AffOrConstraintT = TypeVar("AffOrConstraintT", _isl.Aff, _isl.Constraint)
51
+ AffLikeT = TypeVar("AffLikeT", _isl.Aff, _isl.PwAff)
52
+ ExprLike: TypeAlias = _isl.Aff | _isl.PwAff | _isl.QPolynomial | _isl.PwQPolynomial
53
+ ExprLikeT = TypeVar("ExprLikeT", _isl.Aff, _isl.PwAff,
54
+ _isl.QPolynomial, _isl.PwQPolynomial
55
+ )
56
+ SetLikeT = TypeVar("SetLikeT", bound=_isl.BasicSet | _isl.Set)
57
+ SetOrMap: TypeAlias = _isl.BasicSet | _isl.Set | _isl.BasicMap | _isl.Map
58
+ SetOrMapT = TypeVar("SetOrMapT", bound=SetOrMap)
59
+
60
+ HasSpace: TypeAlias = (
61
+ _isl.Space
62
+ | _isl.Aff
63
+ | _isl.BasicMap
64
+ | _isl.BasicSet
65
+ | _isl.Constraint
66
+ | _isl.LocalSpace
67
+ | _isl.Map
68
+ | _isl.MultiAff
69
+ | _isl.MultiId
70
+ | _isl.MultiPwAff
71
+ | _isl.MultiUnionPwAff
72
+ | _isl.MultiVal
73
+ | _isl.Point
74
+ | _isl.PwAff
75
+ | _isl.PwMultiAff
76
+ | _isl.PwQPolynomial
77
+ | _isl.PwQPolynomialFold
78
+ | _isl.QPolynomial
79
+ | _isl.QPolynomialFold
80
+ | _isl.Set
81
+ | _isl.UnionMap
82
+ | _isl.UnionPwAff
83
+ | _isl.UnionPwMultiAff
84
+ | _isl.UnionPwQPolynomial
85
+ | _isl.UnionPwQPolynomialFold
86
+ | _isl.UnionSet
87
+ )
88
+
89
+
90
+ class IslObject(Protocol):
91
+ def get_ctx(self) -> _isl.Context:
92
+ ...
93
+
94
+ def _wraps_same_instance_as(self, other: object) -> bool:
95
+ ...
96
+
97
+ _base_name: ClassVar[str]
98
+
99
+ # }}}
100
+
101
+
102
+ # {{{ copied verbatim from pytools to avoid numpy/pytools dependency
103
+
104
+ F = TypeVar("F", bound=Callable[..., Any])
105
+
106
+
107
+ class _HasKwargs:
108
+ pass
109
+
110
+
111
+ def _memoize_on_first_arg(function: F, cache_dict_name: str | None = None) -> F:
112
+ """Like :func:`memoize_method`, but for functions that take the object
113
+ in which do memoization information is stored as first argument.
114
+
115
+ Supports cache deletion via ``function_name.clear_cache(self)``.
116
+ """
117
+ from sys import intern
118
+
119
+ if cache_dict_name is None:
120
+ cache_dict_name = intern(
121
+ f"_memoize_dic_{function.__module__}{function.__name__}"
122
+ )
123
+
124
+ def wrapper(obj, *args, **kwargs):
125
+ if kwargs:
126
+ key = (_HasKwargs, frozenset(kwargs.items()), *args)
127
+ else:
128
+ key = args
129
+
130
+ try:
131
+ return getattr(obj, cache_dict_name)[key]
132
+ except AttributeError:
133
+ attribute_error = True
134
+ except KeyError:
135
+ attribute_error = False
136
+
137
+ result = function(obj, *args, **kwargs)
138
+ if attribute_error:
139
+ object.__setattr__(obj, cache_dict_name, {key: result})
140
+ return result
141
+ else:
142
+ getattr(obj, cache_dict_name)[key] = result
143
+ return result
144
+
145
+ def clear_cache(obj):
146
+ object.__delattr__(obj, cache_dict_name)
147
+
148
+ from functools import update_wrapper
149
+ new_wrapper = update_wrapper(wrapper, function)
150
+
151
+ # type-ignore because mypy has a point here, stuffing random attributes
152
+ # into the function's dict is moderately sketchy.
153
+ new_wrapper.clear_cache = clear_cache # type: ignore[attr-defined]
154
+
155
+ return cast("F", new_wrapper)
156
+
157
+ # }}}
158
+
159
+
160
+ def _read_from_str_wrapper(cls, context, s, dims_with_apostrophes):
161
+ """A callable to reconstitute instances from strings for the benefit
162
+ of Python's ``__reduce__`` protocol.
163
+ """
164
+ cls_from_str = cls.read_from_str(context, s)
165
+
166
+ # Apostrophes in dim names have been lost, put them back
167
+ for dim_name, (dt, dim_idx) in dims_with_apostrophes.items():
168
+ cls_from_str = cls_from_str.set_dim_name(dt, dim_idx, dim_name)
169
+
170
+ return cls_from_str
171
+
172
+
173
+ def dim_type_reduce(self: _isl.dim_type):
174
+ return (_isl.dim_type, (int(self),))
175
+
176
+
177
+ def context_reduce(self: _isl.Context):
178
+ from islpy import DEFAULT_CONTEXT, _get_default_context
179
+ if self._wraps_same_instance_as(DEFAULT_CONTEXT):
180
+ return (_get_default_context, ())
181
+ else:
182
+ return (_isl.Context, ())
183
+
184
+
185
+ def context_eq(self: IslObject, other: object):
186
+ return isinstance(other, _isl.Context) and self._wraps_same_instance_as(other)
187
+
188
+
189
+ def context_ne(self: object, other: object) -> bool:
190
+ return not self.__eq__(other)
191
+
192
+
193
+ def generic_reduce(self: IslObject):
194
+ ctx = self.get_ctx()
195
+ prn = _isl.Printer.to_str(ctx)
196
+ prn = getattr(prn, f"print_{self._base_name}")(self)
197
+
198
+ # Reconstructing from string will remove apostrophes in dim names,
199
+ # so keep track of dim names with apostrophes
200
+ dims_with_apostrophes = {
201
+ dname: pos for dname, pos in self.get_var_dict().items()
202
+ if "'" in dname}
203
+
204
+ return (
205
+ _read_from_str_wrapper,
206
+ (type(self), ctx, prn.get_str(), dims_with_apostrophes))
207
+
208
+
209
+ def generic_str(self: IslObject) -> str:
210
+ prn = _isl.Printer.to_str(self.get_ctx())
211
+ getattr(prn, f"print_{self._base_name}")(self)
212
+ return prn.get_str()
213
+
214
+
215
+ def generic_repr(self: IslObject) -> str:
216
+ prn = _isl.Printer.to_str(self.get_ctx())
217
+ getattr(prn, f"print_{self._base_name}")(self)
218
+ return f'{type(self).__name__}("{prn.get_str()}")'
219
+
220
+
221
+ def space_get_id_dict(
222
+ self: _isl.Space,
223
+ dimtype: _isl.dim_type | None = None
224
+ ) -> Mapping[_isl.Id, tuple[_isl.dim_type, int]]:
225
+ """Return a dictionary mapping variable :class:`Id` instances to tuples
226
+ of (:class:`dim_type`, index).
227
+
228
+ :param dimtype: None to get all variables, otherwise
229
+ one of :class:`dim_type`.
230
+ """
231
+ result = {}
232
+
233
+ def set_dim_id(name, tp, idx):
234
+ if name in result:
235
+ raise RuntimeError(f"non-unique var id '{name}' encountered")
236
+ result[name] = tp, idx
237
+
238
+ if dimtype is None:
239
+ types = _CHECK_DIM_TYPES
240
+ else:
241
+ types = [dimtype]
242
+
243
+ for tp in types:
244
+ for i in range(self.dim(tp)):
245
+ name = self.get_dim_id(tp, i)
246
+ if name is not None:
247
+ set_dim_id(name, tp, i)
248
+
249
+ return result
250
+
251
+
252
+ def space_get_var_dict(
253
+ self: _isl.Space,
254
+ dimtype: _isl.dim_type | None = None,
255
+ ignore_out: bool = False
256
+ ) -> Mapping[str, tuple[_isl.dim_type, int]]:
257
+ """Return a dictionary mapping variable names to tuples of
258
+ (:class:`dim_type`, index).
259
+
260
+ :param dimtype: None to get all variables, otherwise
261
+ one of :class:`dim_type`.
262
+ """
263
+ result: dict[str, tuple[_isl.dim_type, int]] = {}
264
+
265
+ def set_dim_name(name: str, tp: _isl.dim_type, idx: int):
266
+ if name in result:
267
+ raise RuntimeError(f"non-unique var name '{name}' encountered")
268
+ result[name] = tp, idx
269
+
270
+ if dimtype is None:
271
+ types = list(_CHECK_DIM_TYPES)
272
+ if ignore_out:
273
+ types = types[:]
274
+ types.remove(_isl.dim_type.out)
275
+ else:
276
+ types = [dimtype]
277
+
278
+ for tp in types:
279
+ for i in range(self.dim(tp)):
280
+ name = self.get_dim_name(tp, i)
281
+ if name is not None:
282
+ set_dim_name(name, tp, i)
283
+
284
+ return result
285
+
286
+
287
+ def space_create_from_names(
288
+ ctx: _isl.Context,
289
+ set: Sequence[str] | None = None,
290
+ in_: Sequence[str] | None = None,
291
+ out: Sequence[str] | None = None,
292
+ params: Sequence[str] = ()
293
+ ) -> _isl.Space:
294
+ """Create a :class:`Space` from lists of variable names.
295
+
296
+ :param set_: names of `set`-type variables.
297
+ :param in_: names of `in`-type variables.
298
+ :param out: names of `out`-type variables.
299
+ :param params: names of parameter-type variables.
300
+ """
301
+ dt = _isl.dim_type
302
+
303
+ if set is not None:
304
+ if in_ is not None or out is not None:
305
+ raise RuntimeError("must pass only one of set / (in_,out)")
306
+
307
+ result = _isl.Space.set_alloc(ctx, nparam=len(params),
308
+ dim=len(set))
309
+
310
+ for i, name in enumerate(set):
311
+ result = result.set_dim_name(dt.set, i, name)
312
+
313
+ elif in_ is not None and out is not None:
314
+ if set is not None:
315
+ raise RuntimeError("must pass only one of set / (in_,out)")
316
+
317
+ result = _isl.Space.alloc(ctx, nparam=len(params),
318
+ n_in=len(in_), n_out=len(out))
319
+
320
+ for i, name in enumerate(in_):
321
+ result = result.set_dim_name(dt.in_, i, name)
322
+
323
+ for i, name in enumerate(out):
324
+ result = result.set_dim_name(dt.out, i, name)
325
+ else:
326
+ raise RuntimeError("invalid parameter combination")
327
+
328
+ for i, name in enumerate(params):
329
+ result = result.set_dim_name(dt.param, i, name)
330
+
331
+ return result
332
+
333
+
334
+ def obj_or(self: SetOrMapT, other: object) -> SetOrMapT:
335
+ try:
336
+ return self.union(other)
337
+ except TypeError:
338
+ return NotImplemented
339
+
340
+
341
+ def obj_and(self: SetOrMapT, other: object) -> SetOrMapT:
342
+ try:
343
+ return self.intersect(other)
344
+ except TypeError:
345
+ return NotImplemented
346
+
347
+
348
+ def obj_sub(self: SetOrMapT, other: object) -> SetOrMapT:
349
+ try:
350
+ return self.subtract(other)
351
+ except TypeError:
352
+ return NotImplemented
353
+
354
+
355
+ def obj_set_coefficients(
356
+ self: AffOrConstraintT,
357
+ dim_tp: _isl.dim_type,
358
+ args: Sequence[_isl.Val | int],
359
+ ) -> AffOrConstraintT:
360
+ """
361
+ :param dim_tp: :class:`dim_type`
362
+ :param args: :class:`list` of coefficients, for indices `0..len(args)-1`.
363
+
364
+ .. versionchanged:: 2011.3
365
+ New for :class:`Aff`
366
+ """
367
+ for i, coeff in enumerate(args):
368
+ self = self.set_coefficient_val(dim_tp, i, coeff)
369
+
370
+ return self
371
+
372
+
373
+ def obj_set_coefficients_by_name(
374
+ self: AffOrConstraintT,
375
+ iterable: Iterable[tuple[str | Literal[1], _isl.Val | int]]
376
+ | Mapping[str | Literal[1], _isl.Val | int],
377
+ name_to_dim: Mapping[str, tuple[_isl.dim_type, int]] | None = None,
378
+ ) -> AffOrConstraintT:
379
+ """Set the coefficients and the constant.
380
+
381
+ :param iterable: a :class:`dict` or iterable of :class:`tuple`
382
+ instances mapping variable names to their coefficients.
383
+ The constant is set to the value of the key '1'.
384
+
385
+ .. versionchanged:: 2011.3
386
+ New for :class:`Aff`
387
+ """
388
+ try:
389
+ coeff_iterable: Iterable[tuple[str | Literal[1], _isl.Val | int]] = \
390
+ list(iterable.items())
391
+ except AttributeError:
392
+ coeff_iterable = \
393
+ cast("Iterable[tuple[str | Literal[1], _isl.Val | int]]", iterable)
394
+
395
+ if name_to_dim is None:
396
+ name_to_dim = obj_get_var_dict(self)
397
+
398
+ for name, coeff in coeff_iterable:
399
+ if name == 1:
400
+ self = self.set_constant_val(coeff)
401
+ else:
402
+ assert name
403
+ tp, idx = name_to_dim[name]
404
+ self = self.set_coefficient_val(tp, idx, coeff)
405
+
406
+ return self
407
+
408
+
409
+ def obj_get_coefficients_by_name(
410
+ self: _isl.Constraint | _isl.Aff,
411
+ dimtype: _isl.dim_type | None = None,
412
+ dim_to_name: Mapping[tuple[_isl.dim_type, int], str] | None = None,
413
+ ) -> dict[str | Literal[1], _isl.Val]:
414
+ """Return a dictionary mapping variable names to coefficients.
415
+
416
+ :param dimtype: None to get all variables, otherwise
417
+ one of :class:`dim_type`.
418
+
419
+ .. versionchanged:: 2011.3
420
+ New for :class:`Aff`
421
+ """
422
+ if dimtype is None:
423
+ types: Sequence[_isl.dim_type] = _CHECK_DIM_TYPES
424
+ else:
425
+ types = [dimtype]
426
+
427
+ result: dict[Literal[1] | str, _isl.Val] = {}
428
+ for tp in types:
429
+ for i in range(self.get_space().dim(tp)):
430
+ coeff = self.get_coefficient_val(tp, i)
431
+ if coeff:
432
+ if dim_to_name is None:
433
+ name = self.get_dim_name(tp, i)
434
+ assert name
435
+ else:
436
+ name = dim_to_name[tp, i]
437
+
438
+ result[name] = coeff
439
+
440
+ const = self.get_constant_val()
441
+ if const:
442
+ result[1] = const
443
+
444
+ return result
445
+
446
+
447
+ def eq_from_names(
448
+ space: _isl.Space,
449
+ coefficients: Mapping[str | Literal[1], _isl.Val | int] | None = None
450
+ ) -> _isl.Constraint:
451
+ """Create a constraint `const + coeff_1*var_1 +... == 0`.
452
+
453
+ :param space: :class:`Space`
454
+ :param coefficients: a :class:`dict` or iterable of :class:`tuple`
455
+ instances mapping variable names to their coefficients
456
+ The constant is set to the value of the key '1'.
457
+
458
+ .. versionchanged:: 2011.3
459
+ Eliminated the separate *const* parameter.
460
+ """
461
+ if coefficients is None:
462
+ coefficients = {}
463
+ c = _isl.Constraint.equality_alloc(space)
464
+ return obj_set_coefficients_by_name(c, coefficients)
465
+
466
+
467
+ def ineq_from_names(
468
+ space: _isl.Space,
469
+ coefficients: Mapping[str | Literal[1], _isl.Val | int] | None = None
470
+ ) -> _isl.Constraint:
471
+ """Create a constraint `const + coeff_1*var_1 +... >= 0`.
472
+
473
+ :param space: :class:`Space`
474
+ :param coefficients: a :class:`dict` or iterable of :class:`tuple`
475
+ instances mapping variable names to their coefficients
476
+ The constant is set to the value of the key '1'.
477
+
478
+ .. versionchanged:: 2011.3
479
+ Eliminated the separate *const* parameter.
480
+ """
481
+ if coefficients is None:
482
+ coefficients = {}
483
+ c = _isl.Constraint.inequality_alloc(space)
484
+ return obj_set_coefficients_by_name(c, coefficients)
485
+
486
+
487
+ def basic_obj_get_constraints(
488
+ self: _isl.BasicSet | _isl.BasicMap
489
+ ) -> list[_isl.Constraint]:
490
+ """Get a list of constraints."""
491
+ result: list[_isl.Constraint] = []
492
+ self.foreach_constraint(result.append)
493
+ return result
494
+
495
+
496
+ def set_get_basic_sets(self: _isl.Set | _isl.BasicSet) -> list[_isl.BasicSet]:
497
+ """Get the list of :class:`BasicSet` instances in this :class:`Set`."""
498
+ result: list[_isl.BasicSet] = []
499
+ self.foreach_basic_set(result.append)
500
+ return result
501
+
502
+
503
+ def map_get_basic_maps(self: _isl.Map) -> list[_isl.BasicMap]:
504
+ """Get the list of :class:`BasicMap` instances in this :class:`Map`."""
505
+ result: list[_isl.BasicMap] = []
506
+ self.foreach_basic_map(result.append)
507
+ return result
508
+
509
+
510
+ def obj_get_id_dict(
511
+ self: HasSpace,
512
+ dimtype: _isl.dim_type | None = None
513
+ ) -> Mapping[_isl.Id, tuple[_isl.dim_type, int]]:
514
+ """Return a dictionary mapping :class:`Id` instances to tuples of
515
+ (:class:`dim_type`, index).
516
+
517
+ :param dimtype: None to get all variables, otherwise
518
+ one of :class:`dim_type`.
519
+ """
520
+ return self.get_space().get_id_dict(dimtype)
521
+
522
+
523
+ @_memoize_on_first_arg
524
+ def obj_get_var_dict(
525
+ self: HasSpace,
526
+ dimtype: _isl.dim_type | None = None
527
+ ) -> Mapping[str, tuple[_isl.dim_type, int]]:
528
+ """Return a dictionary mapping variable names to tuples of
529
+ (:class:`dim_type`, index).
530
+
531
+ :param dimtype: None to get all variables, otherwise
532
+ one of :class:`dim_type`.
533
+ """
534
+ return self.get_space().get_var_dict(
535
+ dimtype, ignore_out=isinstance(self, EXPR_CLASSES))
536
+
537
+
538
+ def obj_get_var_ids(
539
+ self: HasSpace,
540
+ dimtype: _isl.dim_type
541
+ ) -> Sequence[str]:
542
+ """Return a list of :class:`Id` instances for :class:`dim_type` *dimtype*."""
543
+ return [self.get_dim_name(dimtype, i) for i in range(self.dim(dimtype))]
544
+
545
+
546
+ @_memoize_on_first_arg
547
+ def obj_get_var_names(self: HasSpace, dimtype: _isl.dim_type) -> Sequence[str]:
548
+ """Return a list of dim names (in order) for :class:`dim_type` *dimtype*."""
549
+ return [self.get_dim_name(dimtype, i)
550
+ for i in range(self.dim(dimtype))]
551
+
552
+
553
+ def pwaff_get_pieces(self: _isl.PwAff | _isl.Aff) -> list[tuple[_isl.Set, _isl.Aff]]:
554
+ result: list[tuple[_isl.Set, _isl.Aff]] = []
555
+
556
+ def append_tuple(s: _isl.Set, v: _isl.Aff):
557
+ result.append((s, v))
558
+
559
+ self.foreach_piece(append_tuple)
560
+ return result
561
+
562
+
563
+ def pwqpolynomial_get_pieces(
564
+ self: _isl.PwQPolynomial
565
+ ) -> list[tuple[_isl.Set, _isl.QPolynomial]]:
566
+ """
567
+ :return: list of (:class:`Set`, :class:`QPolynomial`)
568
+ """
569
+
570
+ result: list[tuple[_isl.Set, _isl.QPolynomial]] = []
571
+
572
+ def append_tuple(s: _isl.Set, v: _isl.QPolynomial):
573
+ result.append((s, v))
574
+
575
+ self.foreach_piece(append_tuple)
576
+ return result
577
+
578
+
579
+ def pw_get_aggregate_domain(self: _isl.PwAff | _isl.PwQPolynomial) -> _isl.Set:
580
+ """
581
+ :return: a :class:`Set` that is the union of the domains of all pieces
582
+ """
583
+
584
+ result = _isl.Set.empty(self.get_domain_space())
585
+ for dom, _ in self.get_pieces():
586
+ result = result.union(cast("_isl.Set", dom))
587
+
588
+ return result
589
+
590
+
591
+ def qpolynomial_get_terms(self: _isl.QPolynomial) -> list[_isl.Term]:
592
+ """Get the list of :class:`Term` instances in this :class:`QPolynomial`."""
593
+ result: list[_isl.Term] = []
594
+ self.foreach_term(result.append)
595
+ return result
596
+
597
+
598
+ def pwqpolynomial_eval_with_dict(
599
+ self: _isl.PwQPolynomial,
600
+ value_dict: Mapping[str, int | _isl.Val]
601
+ ) -> int:
602
+ """Evaluates *self* for the parameters specified by
603
+ *value_dict*, which maps parameter names to their values.
604
+ """
605
+
606
+ pt = _isl.Point.zero(self.space.params())
607
+
608
+ for i in range(self.space.dim(_isl.dim_type.param)):
609
+ par_name = self.space.get_dim_name(_isl.dim_type.param, i)
610
+ assert par_name
611
+ pt = pt.set_coordinate_val(
612
+ _isl.dim_type.param, i, value_dict[par_name])
613
+
614
+ return self.eval(pt).to_python()
615
+
616
+
617
+ def _number_to_expr_like(template: ExprLikeT, num: int | _isl.Val) -> ExprLikeT:
618
+ number_aff = _isl.Aff.zero_on_domain(template.get_domain_space())
619
+ number_aff = number_aff.set_constant_val(num)
620
+
621
+ if isinstance(template, _isl.Aff):
622
+ return number_aff
623
+ if isinstance(template, _isl.QPolynomial):
624
+ return _isl.QPolynomial.from_aff(number_aff)
625
+
626
+ # everything else is piecewise
627
+
628
+ if template.get_pieces():
629
+ number_pw_aff = _isl.PwAff.empty(template.get_space())
630
+ for set, _ in template.get_pieces():
631
+ number_pw_aff = set.indicator_function().cond(
632
+ number_aff, number_pw_aff)
633
+ else:
634
+ number_pw_aff = _isl.PwAff.alloc(
635
+ _isl.Set.universe(template.domain().space),
636
+ number_aff)
637
+
638
+ if isinstance(template, _isl.PwAff):
639
+ return number_pw_aff
640
+
641
+ elif isinstance(template, _isl.PwQPolynomial):
642
+ return _isl.PwQPolynomial.from_pw_aff(number_pw_aff)
643
+
644
+ else:
645
+ raise TypeError("unexpected template type")
646
+
647
+
648
+ def expr_like_add(self: ExprLikeT, other: ExprLikeT | int | _isl.Val) -> ExprLikeT:
649
+ if not isinstance(other, ExprLike):
650
+ other = _number_to_expr_like(self, other)
651
+
652
+ try:
653
+ return self.add(other)
654
+ except TypeError:
655
+ return NotImplemented
656
+
657
+
658
+ def expr_like_sub(self: ExprLikeT, other: ExprLikeT | int | _isl.Val):
659
+ if not isinstance(other, ExprLike):
660
+ other = _number_to_expr_like(self, other)
661
+
662
+ try:
663
+ return self.sub(other)
664
+ except TypeError:
665
+ return NotImplemented
666
+
667
+
668
+ def expr_like_rsub(self: ExprLikeT, other: ExprLikeT | int | _isl.Val) -> ExprLikeT:
669
+ if not isinstance(other, ExprLike):
670
+ other = _number_to_expr_like(self, other)
671
+
672
+ return -self + other
673
+
674
+
675
+ def expr_like_mul(self: ExprLikeT, other: ExprLikeT | int | _isl.Val) -> ExprLikeT:
676
+ if not isinstance(other, ExprLike):
677
+ other = _number_to_expr_like(self, other)
678
+
679
+ try:
680
+ return self.mul(other)
681
+ except TypeError:
682
+ return NotImplemented
683
+
684
+
685
+ def expr_like_floordiv(self: AffLikeT, other: _isl.Val) -> AffLikeT:
686
+ return self.scale_down_val(other).floor()
687
+
688
+
689
+ def val_rsub(self: _isl.Val, other: _isl.Val) -> _isl.Val:
690
+ return -self + other
691
+
692
+
693
+ def val_bool(self: _isl.Val) -> bool:
694
+ return not self.is_zero()
695
+
696
+
697
+ def val_repr(self: _isl.Val) -> str:
698
+ return f'{type(self).__name__}("{self.to_str()}")'
699
+
700
+
701
+ def val_to_python(self: _isl.Val) -> int:
702
+ if not self.is_int():
703
+ raise ValueError("can only convert integer Val to python")
704
+
705
+ return int(self.to_str())
706
+
707
+
708
+ def obj_eq(self: IslObject, other: object) -> bool:
709
+ assert self.get_ctx() == other.get_ctx(), (
710
+ "Equality-comparing two objects from different ISL Contexts "
711
+ "will likely lead to entertaining (but never useful) results. "
712
+ "In particular, Spaces with matching names will no longer be "
713
+ "equal.")
714
+
715
+ return self.is_equal(other)
716
+
717
+
718
+ def obj_ne(self: object, other: object) -> bool:
719
+ return not self.__eq__(other)
720
+
721
+
722
+ for cls in ALL_CLASSES:
723
+ if hasattr(cls, "is_equal"):
724
+ cls.__eq__ = obj_eq
725
+ cls.__ne__ = obj_ne
726
+
727
+
728
+ def obj_lt(self: SetOrMapT, other: SetOrMapT) -> bool:
729
+ return self.is_strict_subset(other)
730
+
731
+
732
+ def obj_le(self: SetOrMapT, other: SetOrMapT) -> bool:
733
+ return self.is_subset(other)
734
+
735
+
736
+ def obj_gt(self: SetOrMapT, other: SetOrMapT) -> bool:
737
+ return other.is_strict_subset(self)
738
+
739
+
740
+ def obj_ge(self: SetOrMapT, other: SetOrMapT) -> bool:
741
+ return other.is_subset(self)
742
+
743
+
744
+ # {{{ project_out_except
745
+
746
+ def obj_project_out_except(
747
+ obj: SetLikeT,
748
+ names: Collection[str],
749
+ types: Collection[_isl.dim_type]
750
+ ) -> SetLikeT:
751
+ """
752
+ :param types: list of :class:`dim_type` determining
753
+ the types of axes to project out
754
+ :param names: names of axes matching the above which
755
+ should be left alone by the projection
756
+
757
+ .. versionadded:: 2011.3
758
+ """
759
+
760
+ for tp in types:
761
+ while True:
762
+ space = obj.get_space()
763
+ var_dict = space.get_var_dict(tp)
764
+
765
+ all_indices = set(range(space.dim(tp)))
766
+ leftover_indices = {var_dict[name][1] for name in names
767
+ if name in var_dict}
768
+ project_indices = all_indices-leftover_indices
769
+ if not project_indices:
770
+ break
771
+
772
+ min_index = min(project_indices)
773
+ count = 1
774
+ while min_index+count in project_indices:
775
+ count += 1
776
+
777
+ obj = obj.project_out(tp, min_index, count)
778
+
779
+ return obj
780
+
781
+ # }}}
782
+
783
+
784
+ # {{{ eliminate_except
785
+
786
+ def obj_eliminate_except(
787
+ obj: SetLikeT,
788
+ names: Collection[str],
789
+ types: Collection[_isl.dim_type]
790
+ ) -> SetLikeT:
791
+ """
792
+ :param types: list of :class:`dim_type` determining
793
+ the types of axes to eliminate
794
+ :param names: names of axes matching the above which
795
+ should be left alone by the eliminate
796
+
797
+ .. versionadded:: 2011.3
798
+ """
799
+
800
+ for tp in types:
801
+ space = obj.get_space()
802
+ var_dict = space.get_var_dict(tp)
803
+ to_eliminate = (
804
+ set(range(space.dim(tp)))
805
+ - {var_dict[name][1] for name in names
806
+ if name in var_dict})
807
+
808
+ while to_eliminate:
809
+ min_index = min(to_eliminate)
810
+ count = 1
811
+ while min_index+count in to_eliminate:
812
+ count += 1
813
+
814
+ obj = obj.eliminate(tp, min_index, count)
815
+
816
+ to_eliminate -= set(range(min_index, min_index+count))
817
+
818
+ return obj
819
+
820
+ # }}}
821
+
822
+
823
+ # {{{ add_constraints
824
+
825
+ def obj_add_constraints(obj: BasicT, constraints: Iterable[_isl.Constraint]) -> BasicT:
826
+ """
827
+ .. versionadded:: 2011.3
828
+ """
829
+
830
+ for cns in constraints:
831
+ obj = obj.add_constraint(cns)
832
+
833
+ return obj
834
+
835
+ # }}}
836
+
837
+
838
+ def _add_functionality() -> None:
839
+ _isl.dim_type.__reduce__ = dim_type_reduce
840
+
841
+ # {{{ Context
842
+
843
+ _isl.Context.__reduce__ = context_reduce
844
+ _isl.Context.__eq__ = context_eq
845
+ _isl.Context.__ne__ = context_ne
846
+
847
+ # }}}
848
+
849
+ # {{{ generic initialization, pickling
850
+
851
+ for cls in ALL_CLASSES:
852
+ if hasattr(cls, "read_from_str"):
853
+ cls.__reduce__ = generic_reduce
854
+
855
+ # }}}
856
+
857
+ # {{{ printing
858
+
859
+ for cls in ALL_CLASSES:
860
+ if (hasattr(cls, "_base_name")
861
+ and hasattr(_isl.Printer, f"print_{cls._base_name}")):
862
+ cls.__str__ = generic_str
863
+ cls.__repr__ = generic_repr
864
+
865
+ if not hasattr(cls, "__hash__"):
866
+ raise AssertionError(f"not hashable: {cls}")
867
+
868
+ # }}}
869
+
870
+ # {{{ Python set-like behavior
871
+
872
+ for cls in [_isl.BasicSet, _isl.BasicMap, _isl.Set, _isl.Map]:
873
+ cls.__or__ = obj_or
874
+ cls.__ror__ = obj_or
875
+ cls.__and__ = obj_and
876
+ cls.__rand__ = obj_and
877
+ cls.__sub__ = obj_sub
878
+
879
+ # }}}
880
+
881
+ # {{{ Space
882
+
883
+ _isl.Space.create_from_names = staticmethod(space_create_from_names)
884
+ _isl.Space.get_var_dict = space_get_var_dict
885
+ _isl.Space.get_id_dict = space_get_id_dict
886
+
887
+ # }}}
888
+
889
+ # {{{ coefficient wrangling
890
+
891
+ for coeff_class in [_isl.Constraint, _isl.Aff]:
892
+ coeff_class.set_coefficients = obj_set_coefficients
893
+ coeff_class.set_coefficients_by_name = obj_set_coefficients_by_name
894
+ coeff_class.get_coefficients_by_name = obj_get_coefficients_by_name
895
+
896
+ # }}}
897
+
898
+ # {{{ Constraint
899
+
900
+ _isl.Constraint.eq_from_names = staticmethod(eq_from_names)
901
+ _isl.Constraint.ineq_from_names = staticmethod(ineq_from_names)
902
+
903
+ # }}}
904
+
905
+ # {{{ BasicSet
906
+
907
+ _isl.BasicSet.get_constraints = basic_obj_get_constraints
908
+
909
+ # }}}
910
+
911
+ # {{{ BasicMap
912
+
913
+ _isl.BasicMap.get_constraints = basic_obj_get_constraints
914
+
915
+ # }}}
916
+
917
+ # {{{ Set
918
+
919
+ _isl.Set.get_basic_sets = set_get_basic_sets
920
+ _isl.BasicSet.get_basic_sets = set_get_basic_sets
921
+
922
+ # }}}
923
+
924
+ # {{{ Map
925
+
926
+ _isl.Map.get_basic_maps = map_get_basic_maps
927
+
928
+ # }}}
929
+
930
+
931
+ # {{{ common functionality
932
+
933
+ for cls in ALL_CLASSES:
934
+ if hasattr(cls, "get_space") and cls is not _isl.Space:
935
+ cls.get_id_dict = obj_get_id_dict
936
+ cls.get_var_dict = obj_get_var_dict
937
+ cls.get_var_ids = obj_get_var_ids
938
+ cls.get_var_names = obj_get_var_names
939
+
940
+ # }}}
941
+
942
+ # {{{ piecewise
943
+
944
+ _isl.PwAff.get_pieces = pwaff_get_pieces
945
+ _isl.Aff.get_pieces = pwaff_get_pieces
946
+ _isl.PwAff.get_aggregate_domain = pw_get_aggregate_domain
947
+
948
+ _isl.PwQPolynomial.get_pieces = pwqpolynomial_get_pieces
949
+ _isl.PwQPolynomial.get_aggregate_domain = pw_get_aggregate_domain
950
+
951
+ # }}}
952
+
953
+ _isl.QPolynomial.get_terms = qpolynomial_get_terms
954
+
955
+ _isl.PwQPolynomial.eval_with_dict = pwqpolynomial_eval_with_dict
956
+
957
+ # {{{ arithmetic
958
+
959
+ for expr_like_class in ARITH_CLASSES:
960
+ expr_like_class.__add__ = expr_like_add
961
+ expr_like_class.__radd__ = expr_like_add
962
+ expr_like_class.__sub__ = expr_like_sub
963
+ expr_like_class.__rsub__ = expr_like_rsub
964
+ expr_like_class.__mul__ = expr_like_mul
965
+ expr_like_class.__rmul__ = expr_like_mul
966
+ expr_like_class.__neg__ = expr_like_class.neg
967
+
968
+ for qpoly_class in [_isl.QPolynomial, _isl.PwQPolynomial]:
969
+ qpoly_class.__pow__ = qpoly_class.pow
970
+
971
+ for aff_class in [_isl.Aff, _isl.PwAff]:
972
+ aff_class.__mod__ = aff_class.mod_val
973
+ aff_class.__floordiv__ = expr_like_floordiv
974
+
975
+ # }}}
976
+
977
+ # {{{ Val
978
+
979
+ val_cls = _isl.Val
980
+
981
+ val_cls.__add__ = val_cls.add
982
+ val_cls.__radd__ = val_cls.add
983
+ val_cls.__sub__ = val_cls.sub
984
+ val_cls.__rsub__ = val_rsub
985
+ val_cls.__mul__ = val_cls.mul
986
+ val_cls.__rmul__ = val_cls.mul
987
+ val_cls.__neg__ = val_cls.neg
988
+ val_cls.__mod__ = val_cls.mod
989
+ val_cls.__bool__ = val_cls.__nonzero__ = val_bool
990
+
991
+ val_cls.__lt__ = val_cls.lt
992
+ val_cls.__gt__ = val_cls.gt
993
+ val_cls.__le__ = val_cls.le
994
+ val_cls.__ge__ = val_cls.ge
995
+ val_cls.__eq__ = val_cls.eq
996
+ val_cls.__ne__ = val_cls.ne
997
+
998
+ val_cls.__repr__ = val_repr
999
+ val_cls.__str__ = val_cls.to_str
1000
+ val_cls.to_python = val_to_python
1001
+
1002
+ # }}}
1003
+
1004
+ # {{{ rich comparisons
1005
+
1006
+ for cls in [_isl.BasicSet, _isl.BasicMap, _isl.Set, _isl.Map]:
1007
+ cls.__lt__ = obj_lt
1008
+ cls.__le__ = obj_le
1009
+ cls.__gt__ = obj_gt
1010
+ cls.__ge__ = obj_ge
1011
+
1012
+ # }}}
1013
+
1014
+ for c in [_isl.BasicSet, _isl.BasicMap, _isl.Set, _isl.Map]:
1015
+ c.project_out_except = obj_project_out_except
1016
+ c.add_constraints = obj_add_constraints
1017
+
1018
+ for c in [_isl.BasicSet, _isl.Set]:
1019
+ c.eliminate_except = obj_eliminate_except
1020
+
1021
+
1022
+ _add_functionality()
1023
+
1024
+
1025
+ P = ParamSpec("P")
1026
+ ResultT = TypeVar("ResultT")
1027
+
1028
+
1029
+ _DOWNCAST_RE = re.compile(
1030
+ r"Downcast from :class:`([A-Za-z]+)` to :class:`([A-Za-z]+)`.")
1031
+
1032
+
1033
+ _TO_METHODS = {
1034
+ "PwAff": "to_pw_aff",
1035
+ "PwMultiAff": "to_pw_multi_aff",
1036
+ "UnionPwAff": "to_union_pw_aff",
1037
+ "UnionPwMultiAff": "to_union_pw_multi_aff",
1038
+ "LocalSpace": "to_local_space",
1039
+ "Set": "to_set",
1040
+ "UnionSet": "to_union_set",
1041
+ "Map": "to_map",
1042
+ "UnionMap": "to_union_map",
1043
+ }
1044
+
1045
+
1046
+ def _depr_downcast_wrapper(
1047
+ f: Callable[Concatenate[object, P], ResultT],
1048
+ ) -> Callable[Concatenate[object, P], ResultT]:
1049
+ doc = f.__doc__
1050
+ assert doc is not None
1051
+ m = _DOWNCAST_RE.search(doc)
1052
+ assert m, doc
1053
+ basic_cls_name = intern(m.group(1))
1054
+ tgt_cls_name = m.group(2)
1055
+
1056
+ tgt_cls = cast("type", getattr(_isl, tgt_cls_name))
1057
+ is_overload = "Overloaded function" in doc
1058
+ msg = (f"{basic_cls_name}.{f.__name__} "
1059
+ f"with implicit conversion of self to {tgt_cls_name} is deprecated "
1060
+ "and will stop working in 2026. "
1061
+ f"Explicitly convert to {tgt_cls_name}, "
1062
+ f"using .{_TO_METHODS[tgt_cls_name]}().")
1063
+
1064
+ if is_overload:
1065
+ def wrapper(self: object, *args: P.args, **kwargs: P.kwargs) -> ResultT:
1066
+ # "Try to" detect bad invocations of, e.g., Set.union, which is
1067
+ # an overload of normal union and UnionSet.union.
1068
+ if (
1069
+ any(isinstance(arg, tgt_cls) for arg in args)
1070
+ or
1071
+ any(isinstance(arg, tgt_cls) for arg in kwargs.values())
1072
+ ):
1073
+ warn(msg, DeprecationWarning, stacklevel=2)
1074
+
1075
+ return f(self, *args, **kwargs)
1076
+ else:
1077
+ def wrapper(self: object, *args: P.args, **kwargs: P.kwargs) -> ResultT:
1078
+ warn(msg, DeprecationWarning, stacklevel=2)
1079
+
1080
+ return f(self, *args, **kwargs)
1081
+ update_wrapper(wrapper, f)
1082
+ return wrapper
1083
+
1084
+
1085
+ def _monkeypatch_self_downcast_deprecation():
1086
+ for cls in ALL_CLASSES:
1087
+ for attr_name in dir(cls):
1088
+ val = cast("object", getattr(cls, attr_name))
1089
+ doc = getattr(val, "__doc__", None)
1090
+ if doc and "\nDowncast from " in doc:
1091
+ setattr(cls, attr_name, _depr_downcast_wrapper(
1092
+ cast("Callable", val), # pyright: ignore[reportMissingTypeArgument]
1093
+ ))
1094
+
1095
+
1096
+ if not os.environ.get("ISLPY_NO_DOWNCAST_DEPRECATION", None):
1097
+ _monkeypatch_self_downcast_deprecation()