Fix parameter alignment

namedisl/core.py

@@ -40,7 +40,7 @@
 import enum
 import re
 from collections.abc import Callable, Collection, Hashable, Iterable, Mapping, Sequence
-from dataclasses import dataclass
+from dataclasses import dataclass, fields
 from functools import cached_property
 from importlib import metadata
 from typing import (
@@ -360,6 +360,9 @@ def align_obj(
     *, allow_cross_dim_type: bool = False,
     obj_larger_than_space_ok: bool = False,
 ) -> NamedIslObjectT:
+    if named_obj.space.order_equals(space):
+        return named_obj
+
     obj = named_obj._obj
     running_name_to_dim_id = dict(named_obj.space.name_to_dim)
 
@@ -393,6 +396,11 @@ def align_obj(
                         else:
                             another_dim_type = DimType.in_
 
+                    if another_dim_type == DimType.param:
+                        # Only promoted dimensions need consistent parameter IDs.
+                        obj = _set_dim_name(
+                            obj, old_dim_id.dim_type, old_dim_id.dim_index, name)
+
                     obj = obj.move_dims(
                         another_dim_type.as_isl(), 0,
                         old_dim_id.dim_type.as_isl(), old_dim_id.dim_index,
@@ -406,6 +414,10 @@ def align_obj(
                     if not allow_cross_dim_type:
                         raise ValueError("moves across dim_types are not allowed")
 
+                    if target_dt == DimType.param:
+                        obj = _set_dim_name(
+                            obj, old_dim_id.dim_type, old_dim_id.dim_index, name)
+
                     obj = obj.move_dims(
                         target_dim_id.dim_type.as_isl(), target_dim_id.dim_index,
                         old_dim_id.dim_type.as_isl(), old_dim_id.dim_index,
@@ -943,6 +955,19 @@ def as_isl(self) -> IslObjectT_co:
         object.__setattr__(self, "_isl_names_ok", True)  # ruff: ignore[unnecessary-dunder-call]
         return res
 
+    @override
+    def __reduce__(self):
+        return (type(self),
+            (
+                # islpy's pickle protocol requires unique raw dimension names.
+                self.as_isl(),
+                *(
+                    getattr(self, f.name)
+                    for f in fields(self)[1:]
+                )
+            )
+        )
+
     def involves_dims(self, names: Collection[str]) -> bool:
         """True if *self* involves any of the given dimensions."""
         for dt, chunks in chunked_dims_by_type(names, self.space.name_to_dim).items():

namedisl/test/test_namedisl.py

@@ -26,14 +26,17 @@
 """
 
 
+import pickle
+from typing import cast
+
 import pytest
 
 import islpy as isl
 
 import namedisl as nisl
 from .utils_for_tests import generate_random_named_set
 from namedisl import to_named
-from namedisl.core import DimType
+from namedisl.core import DimType, Space, align_obj
 
 
 @pytest.mark.parametrize("ndims", [2, 3, 4, 5])
@@ -228,3 +231,31 @@ def test_ticked_names_are_distinct_names() -> None:
     named_set = nisl.make_set(isl.Set.universe(space))
 
     assert named_set.space.names == frozenset({"x", "x'"})
+
+
+def test_align_obj_returns_already_aligned_object() -> None:
+    named_set = nisl.make_set("{ [i] : i >= 0 }")
+
+    assert align_obj(named_set, named_set.space) is named_set
+
+
+def test_align_obj_promotes_renamed_dimension_with_its_named_id() -> None:
+    renamed_set = nisl.make_set("{ [i] : i = 4 }").rename_dims((("i", "n"),))
+    parameter_space = Space.from_names(param=["n"], out=[])
+
+    aligned = align_obj(
+        renamed_set, parameter_space, allow_cross_dim_type=True)
+
+    assert aligned.equals(nisl.make_set("[n] -> { [] : n = 4 }"))
+
+
+def test_pickling_restores_unique_raw_dimension_names() -> None:
+    raw_space = isl.Space.set_alloc(isl.DEFAULT_CONTEXT, 0, 2)
+    named_set = nisl.Set(
+        isl.Set.universe(raw_space), Space.from_names(param=[], out=["i", "j"]))
+
+    restored = cast("nisl.Set", pickle.loads(pickle.dumps(named_set)))
+
+    assert restored.equals(named_set)
+    assert restored._obj.get_dim_name(isl.dim_type.set, 0) == "i"
+    assert restored._obj.get_dim_name(isl.dim_type.set, 1) == "j"