fix(language): at refuses a read that consumes a column over the dimension it joins on, as sum does

src/math_spec/_expression_parser.py

@@ -132,24 +132,23 @@ def __str__(self) -> str:
 
 
 @dataclass(frozen=True)
-class RelationNode:
-    """A resolved ``by=`` — the relation, and the use the call makes of it.
-
-    ``use`` is a :class:`Direction` for ``sum`` and ``at``, and a
-    :class:`Partition` for ``shift`` and ``sum_back``.
-    ``dimensions`` is the fine side — what ``sum`` consumes and ``at``
-    produces — and ``into`` the coarse dims, which ``sum`` produces and ``at``
-    consumes. The roles joined on are the operand's to carry, and the operator
-    passes them through.
-    """
+class DirectionNode:
+    """A resolved ``by=`` on ``sum`` or ``at``: the relation, read in the :class:`Direction` the call names."""
 
-    name: str
-    dimensions: tuple[str, ...]
-    into: tuple[str, ...]
-    use: Direction | Partition
+    direction: Direction
 
     def __str__(self) -> str:
-        return self.name
+        return self.direction.name
+
+
+@dataclass(frozen=True)
+class PartitionNode:
+    """A resolved ``by=`` on ``shift`` or ``sum_back``: the relation, as the :class:`Partition` the call steps inside."""
+
+    partition: Partition
+
+    def __str__(self) -> str:
+        return self.partition.name
 
 
 @dataclass(frozen=True)
@@ -284,7 +283,8 @@ def __str__(self) -> str:
     | ParameterNode
     | DualNode
     | DimensionNode
-    | RelationNode
+    | DirectionNode
+    | PartitionNode
     | EdgeNode
     | KeywordNode
     | UnaryOperatorNode
@@ -338,7 +338,7 @@ def operand(node: ArithmeticNode) -> str:
 #: ``sum(x, along=d)``, ``sum(x, by=l)``, ``shift(..., edge='wrap')``. None of
 #: the three is data, so none may stand in arithmetic — which is why the passes
 #: that walk a value position refuse them together.
-KwargNode = DimensionNode | RelationNode | EdgeNode
+KwargNode = DimensionNode | DirectionNode | PartitionNode | EdgeNode
 
 #: What resolution rewrites away: a bare name, whose kind only the schema
 #: knows, and the two kwarg-only literals its kwarg consumes. Meeting one

src/math_spec/advice.py

@@ -69,12 +69,11 @@ def _never_an_axis(program: Program) -> list[Advice]:
 def _produced_axes(program: Program) -> set[str]:
     """The axes the expressions create beyond what any declaration indexes.
 
-    ``sum(by=)`` lands on its target and ``at()`` spreads onto its fine dimension.
+    ``sum(by=)`` lands on its target and ``at()`` spreads onto its fine
+    dimension: either way, the dims the direction produces.
     """
     axes: set[str] = set()
     for node in walk(*program.expressions):
-        if isinstance(node, GroupSum):
-            axes.update(node.into)
-        elif isinstance(node, At):
-            axes.update(node.over)
+        if isinstance(node, GroupSum | At):
+            axes.update(node.direction.produced_dims)
     return axes

src/math_spec/dimensions.py

@@ -23,14 +23,15 @@
     ComparisonNode,
     DefinitionNode,
     DimensionNode,
+    DirectionNode,
     DualNode,
     EdgeNode,
     FunctionCallNode,
     KwargNode,
     NumberNode,
     ParameterNode,
     ParsedNode,
-    RelationNode,
+    PartitionNode,
     UnaryOperatorNode,
     UnresolvedNode,
     VariableNode,
@@ -42,6 +43,7 @@
 from math_spec.program import (
     DimensionComparisonNode,
     DimensionPositionNode,
+    Direction,
     Mask,
     ParameterComparisonNode,
     ParameterDefinedNode,
@@ -154,36 +156,33 @@ def _sum_dims(node: FunctionCallNode, inner: frozenset[str], schema: Spec, conte
             )
         return inner - {consumed.name}
 
-    assert isinstance(by, RelationNode)
-    if missing := sorted(set(by.dimensions) - inner):
+    assert isinstance(by, DirectionNode), 'resolution reads sum(by=) in a direction'
+    direction = by.direction
+    if missing := sorted(set(direction.consumed_dims) - inner):
         raise DimensionError(
             _not_carried(
                 context,
-                f'sum(by={by}) consumes {missing}, the dims it reads from,',
+                f'sum(by={direction.name}) consumes {missing}, the dims it reads from,',
                 inner,
                 'drop the sum, or fix the dim',
             )
         )
-    _check_lands_clear(f'sum(by={by})', set(by.into), set(by.dimensions), inner, context)
-    _check_joined(f'sum(by={by})', by, inner, context)
-    return (inner - set(by.dimensions)) | set(by.into)
+    return _read_dims(f'sum(by={direction.name})', direction, inner, context)
 
 
 def _at_dims(node: FunctionCallNode, inner: frozenset[str], schema: Spec, context: str) -> frozenset[str]:
     """``at`` is the adjoint of ``sum(by=)``: it consumes the dims a sum produces and produces the ones it consumes."""
     by = node.kwargs['by']
-    assert isinstance(by, RelationNode)
-    absent = sorted(set(by.into) - inner)
-    if absent:
+    assert isinstance(by, DirectionNode), 'resolution reads at(by=) in a direction'
+    direction = by.direction
+    if absent := sorted(set(direction.consumed_dims) - inner):
         raise DimensionError(
-            f'{context}: at(by={by}) reads through '
+            f'{context}: at(by={direction.name}) reads through '
             f'{absent}, which the expression does not carry (dims '
             f'{sorted(inner)}). A pullback needs the coarse dims to read *from* — '
             f'sum is the direction that produces them.'
         )
-    _check_lands_clear(f'at(by={by})', set(by.dimensions), set(by.into), inner, context)
-    _check_joined(f'at(by={by})', by, inner, context)
-    return (inner - set(by.into)) | set(by.dimensions)
+    return _read_dims(f'at(by={direction.name})', direction, inner, context)
 
 
 def _translation_dims(node: FunctionCallNode, inner: frozenset[str], schema: Spec, context: str) -> frozenset[str]:
@@ -202,34 +201,42 @@ def _translation_dims(node: FunctionCallNode, inner: frozenset[str], schema: Spe
     _check_named_amount(node, over.name, inner, schema, context)
     _check_amount_form(node, context)
     _check_edge(node, context)
-    partition = node.kwargs.get('by')
-    if partition is not None:
-        assert isinstance(partition, RelationNode)
-        _check_joined(f'{node.name}(along={over.name}, by={partition})', partition, inner, context)
+    by = node.kwargs.get('by')
+    if by is not None:
+        assert isinstance(by, PartitionNode), "resolution reads a translation's by= as a partition"
+        _check_joined(f'{node.name}(along={over.name}, by={by.partition.name})', by.partition, inner, context)
     return inner
 
 
-def _check_lands_clear(call: str, produced: set[str], consumed: set[str], inner: frozenset[str], context: str) -> None:
-    """The dims a call lands on are its own to bring, so the operand does not already carry one.
+def _read_dims(call: str, direction: Direction, inner: frozenset[str], context: str) -> frozenset[str]:
+    """The dims after a relation is read in *direction*: the consumed go, the produced arrive, the joined stay.
 
-    Where it does, the call would tie the operand's axis to the one it
-    produces rather than adding it, and it reads the same either way.
-    A relation into its own dimension is not that case: there the dim landed
-    on is the dim just consumed, so every factor is read at the coordinate
-    the sum runs over, and nothing is tied.
+    The dims a call lands on are its own to bring, so the operand does not
+    already carry one. Where it does, the call would tie the operand's axis to
+    the one it produces rather than adding it, and it reads the same either
+    way. A relation into its own dimension is not that case: there the dim
+    landed on is the dim just consumed, so every factor is read at the
+    coordinate the sum runs over, and nothing is tied.
     """
+    consumed, produced = set(direction.consumed_dims), set(direction.produced_dims)
     if clash := sorted((produced & inner) - consumed):
         raise DimensionError(
             f'{context}: {call} lands on {clash}, which the expression already carries.\n'
             f'A call brings the dims it lands on, so that reading it tells you what it '
             f'adds. Move the factor carrying {clash} outside the operator, or read to a column '
             f'over another dimension.'
         )
+    _check_joined(call, direction, inner, context)
+    return (inner - consumed) | produced
+
 
+def _check_joined(call: str, use: Direction | Partition, inner: frozenset[str], context: str) -> None:
+    """The columns a call joins on are read at their dimensions, so the operand carries every one, each once.
 
-def _check_joined(call: str, by: RelationNode, inner: frozenset[str], context: str) -> None:
-    """The columns a call joins on are read at their dimensions, so the operand carries every one, each once."""
-    use = by.use
+    A joined dimension the call also consumes is the same ambiguity as two
+    joined columns over one dimension: the operand's one coordinate would
+    have to be read as both. A partition consumes nothing.
+    """
     dims = use.joined_dims
     if missing := sorted(set(dims) - inner):
         raise DimensionError(
@@ -238,7 +245,8 @@ def _check_joined(call: str, by: RelationNode, inner: frozenset[str], context: s
             f'read between two of its columns and joined at the others — index the operand by them, or '
             f'read it between different columns.'
         )
-    if twice := sorted({d for d in dims if dims.count(d) > 1 or d in by.dimensions}):
+    consumed = use.consumed_dims if isinstance(use, Direction) else ()
+    if twice := sorted({d for d in dims if dims.count(d) > 1 or d in consumed}):
         raise DimensionError(
             f"{context}: {call} joins '{use.name}' on {twice} through more than one column, and the operand "
             f'carries each dimension once. Read between different columns, or use a relation whose joined '
@@ -421,9 +429,10 @@ def _check_named_amount(node: FunctionCallNode, over: str, inner: frozenset[str]
             f'carries it. A named {words.noun} that varies over the axis it steps along is {words.varies} '
             f"— declare '{amount.name}' over dims '{over}' is not one of."
         )
-    partition = node.kwargs.get('by')
-    use = partition.use if isinstance(partition, RelationNode) else None
-    groups = frozenset(use.dim(v) for v in use.group) if isinstance(use, Partition) else frozenset()
+    by = node.kwargs.get('by')
+    groups = (
+        frozenset(by.partition.dim(v) for v in by.partition.group) if isinstance(by, PartitionNode) else frozenset()
+    )
     if stray := sorted(frozenset(declared.dims) - inner - groups):
         raise DimensionError(
             f'{context}: {node.name}({kwarg}={amount.name}) reads its {words.noun} at the coordinate it '

src/math_spec/lowering.py

@@ -22,13 +22,14 @@
     CasesNode,
     DefinitionNode,
     DimensionNode,
+    DirectionNode,
     DualNode,
     EdgeNode,
     FunctionCallNode,
     KwargNode,
     NumberNode,
     ParameterNode,
-    RelationNode,
+    PartitionNode,
     UnaryOperatorNode,
     UnresolvedNode,
     VariableNode,
@@ -141,11 +142,7 @@ def lower_program(expanded: _ExpandedSpec) -> program.Program:
 
     dimensions = {
         dname: program.DimensionDeclaration(
-            tuple(
-                program.RelationDeclaration(lname, lk.pairs, lk.key_roles)
-                for lname, lk in expanded.relations.items()
-                if dname in lk.dims
-            ),
+            tuple(lk for lk in resolved.relations.values() if dname in lk.dims),
             ddef.dtype,
         )
         for dname, ddef in expanded.dimensions.items()
@@ -272,16 +269,14 @@ def sum(self, node: FunctionCallNode) -> program.ExpressionNode:
             consumed = node.kwargs['over']
             assert isinstance(consumed, DimensionNode), 'resolution refuses a over= that is not a dimension'
             return program.Sum(operand, (consumed.name,))
-        assert isinstance(by_node, RelationNode), 'resolution refuses a by= that is not a relation'
-        assert isinstance(by_node.use, program.Direction), 'resolution reads sum(by=) in a direction'
-        return program.GroupSum(operand, direction=by_node.use)
+        assert isinstance(by_node, DirectionNode), 'resolution reads sum(by=) in a direction'
+        return program.GroupSum(operand, direction=by_node.direction)
 
     def at(self, node: FunctionCallNode) -> program.ExpressionNode:
         """``at(x, by=relation)`` — the adjoint of :meth:`sum`'s ``by=`` form."""
         by_node = node.kwargs['by']
-        assert isinstance(by_node, RelationNode), 'resolution refuses a by= that is not a relation'
-        assert isinstance(by_node.use, program.Direction), 'resolution reads at(by=) in a direction'
-        return program.At(self.expr(node.args[0]), direction=by_node.use)
+        assert isinstance(by_node, DirectionNode), 'resolution reads at(by=) in a direction'
+        return program.At(self.expr(node.args[0]), direction=by_node.direction)
 
     def sum_back(self, node: FunctionCallNode) -> program.ExpressionNode:
         """``sum_back(x, along=d, window=w)`` — a trailing window along one dimension.
@@ -353,9 +348,8 @@ def _partition_of(node: FunctionCallNode) -> program.Partition | None:
     by_node = node.kwargs.get('by')
     if by_node is None:
         return None
-    assert isinstance(by_node, RelationNode)
-    assert isinstance(by_node.use, program.Partition), "resolution reads a translation's by= as a partition"
-    return by_node.use
+    assert isinstance(by_node, PartitionNode), "resolution reads a translation's by= as a partition"
+    return by_node.partition
 
 
 def _bound_expression(value: float | str) -> program.ExpressionNode: