feat(program): each expression node states the dimensions it is read over

CHANGELOG.md

@@ -12,6 +12,7 @@ it releases that version ([RELEASING.md](https://github.com/energy-models/mathsp
 
 ## Upcoming version
 
+- feat(program): each expression node states the dimensions it is read over ([#899](https://github.com/energy-models/mathspec/pull/899))
 - fix(language): two files that declare their dimensions in different orders load to the same program by sorting dimension lists by name ([#911](https://github.com/energy-models/mathspec/pull/911))
 - refactor(language)!: a file names a condition under `conditions:` rather than `masks:` ([#907](https://github.com/energy-models/mathspec/pull/907))
 - refactor(api)!: the docs, the error messages and the program classes call a true-or-false `where` expression a condition rather than a predicate ([#904](https://github.com/energy-models/mathspec/pull/904))

docs/reference/language/errors.md

@@ -57,6 +57,13 @@ variable with no constraint row.
 | `SchemaError`    | `LanguageError` | Something in the file: an unknown key, a malformed entry, or a bad symbol table                                                 |
 | `DimensionError` | `LanguageError` | Dimensions that disagree, such as a constraint whose expression does not equal its `dims`                                       |
 
+A load reads every expression and `where:` in the file, and then lists each
+fault it found, one to a line. The error is a `DimensionError` when every fault
+breaks a dimension rule, and a `SchemaError` when any fault is of another kind.
+After that, the load checks the rules on whole entries, such as a constraint
+whose expression is not indexed over exactly its `dims:`. That `DimensionError` names the
+first entry that breaks a rule.
+
 You can reproduce each of these errors from the YAML alone.
 
 ## What the language will not express

docs/reference/reading.md

@@ -147,6 +147,19 @@ entry's expression, the same object that `program.expressions[name].expression`
 holds, and its value is the body's value. `children()` steps into the body, so
 a walk reads through it.
 
+Every expression node has `dims`: the dimensions over which it is indexed,
+sorted by name. A sum has the dimensions its operand has,
+less the ones it sums over. A `Join` under a sum through a relation has the
+dimensions of that sum. A `Join` also opens one `Axis` for each column it
+drops, but an axis is not a dimension and is not in `dims`. An
+`ExpressionReference` has the dimensions of its body.
+
+A node's `dims` are always sorted by name. The order of `dimensions:` and the
+order of an entry's `dims:` do not change them. An entry that declares
+`dims: [snapshot, generator]` has
+`program.expressions[name].dims == ('snapshot', 'generator')`, and each
+`ExpressionReference` to it has `dims == ('generator', 'snapshot')`.
+
 A `ConditionReference` stands where a `conditions:` entry is read. Its `body` is the
 condition of the entry, the root of `program.conditions[name].condition`, and it is true
 where the body is true. `condition_children()` steps into the body, so `.atoms`,

src/mathspec/_condition_resolver.py

@@ -25,7 +25,7 @@
     UnresolvedConditionNode,
     UnresolvedCountNode,
 )
-from mathspec._dimensions import dims_of, join_dims
+from mathspec._dimensions import join_dims
 from mathspec._expansion import expand
 from mathspec._expression_parser import (
     ArithmeticNode,
@@ -266,7 +266,7 @@ def _joined(self, node: UnresolvedConditionCallNode, mask: Mask) -> Condition |
         try:
             dims = join_dims(by, mask.dims, context, 'the condition')
         except DimensionError as refusal:
-            self.errors.append(str(refusal))
+            self.errors.append(self.ns.refused_dims(refusal))
             return node
         return JoinedCondition(mask, by, tuple(sorted(dims)))
 
@@ -408,7 +408,7 @@ def _expression_comparison(self, node: UnresolvedComparisonNode) -> ExpressionCo
             else:
                 try:
                     degree.check_expression(side, context)
-                    dims |= dims_of(side, ns.spec, context)
+                    dims |= set(side.dims)
                 except LanguageError as e:
                     self.errors.append(str(e))
         if len(self.errors) > found:

src/mathspec/_dimensions.py

@@ -4,10 +4,12 @@
 
 """Static dim-set checking — a type system whose type is a set of dim names.
 
-Every node's dim set is computable before any data is attached, so this pass runs
-at load on the resolved tree. The per-node rules are the "Dim algebra" table in
-``docs/reference/language/expressions.md``; a constraint's two sides together
-must equal its ``dims``, and a condition or a bound may not exceed the frame.
+Every node states its dims, so every rule is decided at load. Resolution holds
+each node to the rule of its operator as it builds the node ([`check_node`][]);
+the rules are the "Dim algebra" table in
+``docs/reference/language/expressions.md``. [`check_spec`][] holds each entry
+to its own: a constraint's two sides together are indexed over exactly its
+``dims``, and a where or a bound is indexed over no dim the entry is not.
 """
 
 from __future__ import annotations
@@ -16,115 +18,73 @@
 
 from mathspec.errors import DimensionError, case_context
 from mathspec.program import (
-    Add,
     Cases,
-    Constant,
     CountComparison,
     DimensionComparison,
     DimensionPosition,
-    Divide,
-    Dual,
     Expression,
     ExpressionComparison,
-    ExpressionReference,
     Join,
     JoinColumns,
     JoinedCondition,
     Mask,
-    Multiply,
-    Negate,
     ParameterComparison,
     ParameterDefined,
-    ParameterReference,
     Partition,
-    Power,
     RelationComparison,
     RelationDefined,
     RelationPairComparison,
     Sum,
     Translate,
     TranslatedCondition,
     VariableDefined,
-    VariableReference,
     WindowSum,
-    children,
 )
 
 if TYPE_CHECKING:
+    from collections.abc import Iterable
+
     from mathspec.program import Program
     from mathspec.spec import Spec
 
 
-def dims_of(node: Expression, spec: Spec, context: str) -> frozenset[str]:
-    """The dim set of a resolved expression, checking every rule on the way.
+def check_node(node: Sum | Join | Translate | WindowSum, spec: Spec, context: str) -> None:
+    """The dim rule *node*'s own operator holds it to, on the dims its operand states.
+
+    Resolution calls it on each node it builds, so a tree is checked once,
+    bottom up, and no pass walks it again.
 
     Raises:
-        DimensionError: On the first rule broken.
+        DimensionError: *node* breaks the rule.
     """
-    if isinstance(node, Constant):
-        return frozenset()
-
-    if isinstance(node, ParameterReference):
-        return frozenset({**spec.parameters, **spec.given.parameters}[node.name].dims)
-
-    if isinstance(node, VariableReference):
-        columns = {**spec.variables, **spec.given.variables, **spec.given.expressions}
-        return frozenset(columns[node.name].dims or ())
-
-    if isinstance(node, Dual):
-        return frozenset({**spec.constraints, **spec.given.constraints}[node.constraint].dims)
-
-    if isinstance(node, ExpressionReference):
-        return _named_dims(node, spec, context)
-
-    if isinstance(node, Cases):
-        return frozenset().union(*(dims_of(region.value, spec, context) for region in node.regions))
-
-    if isinstance(node, Negate | Add | Multiply | Power | Divide):
-        return frozenset().union(*(dims_of(child, spec, context) for child in children(node)))
-
+    inner = frozenset(node.operand.dims)
     if isinstance(node, Sum):
-        return _sum_dims(node, spec, context)
-    inner = dims_of(node.operand, spec, context)
-    if isinstance(node, Join):
-        assert not node.columns.axes, 'a join that opens axes is read only by the sum that closes them'
-        return join_dims(node.columns, inner, context, 'the expression')
-    if isinstance(node, Translate | WindowSum):
-        return _translation_dims(node, inner, spec, context)
-
-    assert_never(node)
-
-
-def _named_dims(node: ExpressionReference, spec: Spec, context: str) -> frozenset[str]:
-    """The dims an entry declares where it declares them, else those of its expression."""
-    declared = spec.expressions[node.name].dims
-    if declared is not None:
-        return frozenset(declared)
-    return dims_of(node.body, spec, context)
+        _check_sum(node, inner, context)
+    elif isinstance(node, Join):
+        join_dims(node.columns, inner, context, 'the expression')
+    else:
+        _check_translation(node, inner, spec, context)
 
 
 def _not_carried(context: str, call: str, inner: frozenset[str], rewrite: str) -> str:
     """The refusal for an operator reaching a dim its operand does not carry; *rewrite* is the operator's own."""
     return f'{context}: {call} but the expression has only the dims {sorted(inner)}. {rewrite}.'
 
 
-def _sum_dims(node: Sum, spec: Spec, context: str) -> frozenset[str]:
+def _check_sum(node: Sum, inner: frozenset[str], context: str) -> None:
     """``sum`` reduces each dim in ``over`` away, so the operand carries every one.
 
     The axes a join opens never reach a frame: the sum over the join closes
-    every one, so its frame is the join's frame less them.
+    every one, and the join's own rules are the sum's.
     """
     if isinstance(node.operand, Join) and node.operand.columns.axes:
-        join = node.operand
-        assert node.over == join.columns.axes, 'a sum through a relation closes exactly the axes its join opens'
-        return join_dims(join.columns, dims_of(join.operand, spec, context), context, 'the expression')
-    inner = dims_of(node.operand, spec, context)
+        assert node.over == node.operand.columns.axes, 'a sum through a relation closes exactly the axes its join opens'
+        return
     for summed in node.over:
         if summed.dimension not in inner:
             raise DimensionError(
                 _not_carried(context, f'sum(over={summed})', inner, 'Remove the sum, or correct the dimension')
             )
-    return inner - {axis.dimension for axis in node.over}
 
 
 def join_dims(columns: JoinColumns, inner: frozenset[str], context: str, operand: str) -> frozenset[str]:
@@ -160,14 +120,24 @@ def join_dims(columns: JoinColumns, inner: frozenset[str], context: str, operand
             f'dimension.'
         )
     _check_joined(call, columns, inner, context)
-    return (inner - set(columns.joined_dims)) | set(columns.grouped_dims)
+    return joined(columns, inner)
+
+
+def sorted_dims(*dims: Iterable[str]) -> tuple[str, ...]:
+    """Every dim in *dims*, once, sorted by name: what a node states as ``dims``."""
+    return tuple(sorted(set().union(*dims)))
+
+
+def joined(columns: JoinColumns, inner: Iterable[str]) -> frozenset[str]:
+    """The dims a join on *columns* leaves of *inner*: the dims joined on go, and the dims grouped by arrive."""
+    return (frozenset(inner) - set(columns.joined_dims)) | set(columns.grouped_dims)
 
 
 #: The operator name a file writes for each translation, which its refusals quote.
 _VERBS: dict[type[Translate | WindowSum], str] = {Translate: 'shift', WindowSum: 'sum_back'}
 
 
-def _translation_dims(node: Translate | WindowSum, inner: frozenset[str], spec: Spec, context: str) -> frozenset[str]:
+def _check_translation(node: Translate | WindowSum, inner: frozenset[str], spec: Spec, context: str) -> None:
     """``shift`` and ``sum_back`` keep every dim, and a named amount and a partition are checked here."""
     verb = _VERBS[type(node)]
     if node.along not in inner:
@@ -183,7 +153,6 @@ def _translation_dims(node: Translate | WindowSum, inner: frozenset[str], spec:
     if node.partition is not None:
         within = f'{node.partition.name}[{", ".join(node.partition.grouped)}]'
         _check_joined(f'{verb}(along={node.along}, within={within})', node.partition, inner, context)
-    return inner
 
 
 def _check_joined(call: str, use: JoinColumns | Partition, inner: frozenset[str], context: str) -> None:
@@ -264,15 +233,15 @@ def check_spec(spec: Spec, program: Program) -> None:
             continue
         frame = frozenset(declared.dims)
         if not isinstance(entry.expression, Cases):
-            indexed = _check_body_dims(entry.expression, spec, frame, f"Named expression '{ename}'")
+            indexed = _check_body_dims(entry.expression, frame, f"Named expression '{ename}'")
             _check_dims_are_indexed(ename, frame, indexed, 'the expression is not')
             continue
         indexed = frozenset[str]()
         for region, label in zip(entry.expression.regions, [*declared.cases, None], strict=True):
             context = case_context(ename, label)
             if label is not None:
                 _check_condition_dims(region.when, frame, context)
-            indexed |= _check_value_dims(region.value, spec, frame, context) | region.when.dims
+            indexed |= _check_value_dims(region.value, frame, context) | region.when.dims
         _check_dims_are_indexed(ename, frame, indexed, 'no case value and no `when:` is')
 
     for ename, entry in program.expressions.items():
@@ -289,7 +258,7 @@ def check_spec(spec: Spec, program: Program) -> None:
         frame = frozenset(constraint.dims)
         context = f"Constraint '{cname}'"
         _check_condition_dims(constraint.where, frame, context)
-        got = dims_of(constraint.lhs, spec, context) | dims_of(constraint.rhs, spec, context)
+        got = frozenset((*constraint.lhs.dims, *constraint.rhs.dims))
         if got != frame:
             stray, missing = sorted(got - frame), sorted(frame - got)
             detail = (
@@ -302,7 +271,7 @@ def check_spec(spec: Spec, program: Program) -> None:
 
     if program.objective is not None:
         context = 'The objective'
-        got = dims_of(program.objective.expression, spec, context)
+        got = program.objective.expression.dims
         if got:
             raise DimensionError(
                 f'{context}: the expression carries dims {sorted(got)}, and an objective is one '
@@ -311,13 +280,13 @@ def check_spec(spec: Spec, program: Program) -> None:
             )
 
 
-def _check_value_dims(node: Expression, spec: Spec, frame: frozenset[str], context: str) -> frozenset[str]:
+def _check_value_dims(node: Expression, frame: frozenset[str], context: str) -> frozenset[str]:
     """A region's value may only carry dims the frame does — the ``otherwise:`` included.
 
     A wider one would give the quantity dims its entry does not, which is
     the second answer a ``dims:`` exists to avoid.
     """
-    got = dims_of(node, spec, context)
+    got = frozenset(node.dims)
     if not got <= frame:
         raise DimensionError(
             f'{context}: the value carries dims {sorted(got - frame)} outside the dims: '
@@ -326,9 +295,9 @@ def _check_value_dims(node: Expression, spec: Spec, frame: frozenset[str], conte
     return got
 
 
-def _check_body_dims(node: Expression, spec: Spec, frame: frozenset[str], context: str) -> frozenset[str]:
+def _check_body_dims(node: Expression, frame: frozenset[str], context: str) -> frozenset[str]:
     """A plain entry's expression is indexed only over dims its ``dims:`` gives; the dims it is indexed over come back."""
-    got = dims_of(node, spec, context)
+    got = frozenset(node.dims)
     if not got <= frame:
         raise DimensionError(
             f'{context}: the body carries dims {sorted(got - frame)} outside the dims: {sorted(frame)}. '