diff --git a/mathics/core/expression.py b/mathics/core/expression.py index 0be8e99e5..369aed1a3 100644 --- a/mathics/core/expression.py +++ b/mathics/core/expression.py @@ -371,6 +371,8 @@ class Expression(BaseElement, NumericOperators, EvalMixin): _cache: ExpressionCache | None # (head, elements, element_order) of the last element_order built. _element_order_cache: tuple[BaseElement, tuple[BaseElement, ...], tuple] | None + # Whether no Slot occurs in the expression: see _has_no_slots(). + _slot_free: bool | None = None elements_properties: ElementsProperties | None options: dict[str, Any] | None pattern_sequence: bool @@ -614,6 +616,7 @@ def elements(self, values: Iterable): # Set to build self.elements_properties on next evaluation() self.elements_properties = None self._cache = None + self._slot_free = None self._element_order_cache = None def equal2(self, rhs: Any) -> bool | None: @@ -1732,6 +1735,7 @@ def set_head(self, head: Symbol): raise TypeError("Attempt to turn an Expression into a ListExpression") self._head = head self._cache = None + self._slot_free = None self._element_order_cache = None def set_element(self, index: int, value): @@ -2012,6 +2016,28 @@ def replace_vars(self, vars, options=None, in_function=True) -> "Expression": return self return Expression(head, *new_elements) + def _has_no_slots(self) -> bool: + """ + True if no Slot or SlotSequence occurs in the expression, heads + included. Worked out once and kept with the expression, so that a + part shared by every application of a pure function is looked at + once; changing the head or the elements forgets it. + """ + flag = self._slot_free + if flag is None: + head = self._head + flag = ( + head is not SymbolSlot + and head is not SymbolSlotSequence + and (not isinstance(head, Expression) or head._has_no_slots()) + and all( + not isinstance(element, Expression) or element._has_no_slots() + for element in self._elements + ) + ) + self._slot_free = flag + return flag + def replace_slots(self, slots, evaluation): """ Replaces Slots (#1, ##, etc) by the corresponding values in `slots` @@ -2038,6 +2064,10 @@ def replace_slots(self, slots, evaluation): elif self._head is SymbolFunction and len(self._elements) == 1: # do not replace Slots in nested Functions return self + if self._has_no_slots(): + # No slot anywhere below, e.g. a graph passed to a function and + # written into the body of a pure function: nothing to walk. + return self head = self._head.replace_slots(slots, evaluation) elements = [ element.replace_slots(slots, evaluation) for element in self._elements diff --git a/test/core/test_replace_slots.py b/test/core/test_replace_slots.py new file mode 100644 index 000000000..777b2ed11 --- /dev/null +++ b/test/core/test_replace_slots.py @@ -0,0 +1,68 @@ +# -*- coding: utf-8 -*- +""" +Applying a pure function does not walk the parts of its body that have no +slot in them. +""" + +from test.helper import check_evaluation, session + +import pytest + +from mathics.core.expression import Expression + + +def test_slot_free_part_is_not_walked(monkeypatch): + big = session.evaluate("Table[{i, w -> i}, {i, 300}]") + body = session.evaluate("Hold[{#1, Length[x]}]").elements[0] + # The body {#1, Length[]}, as when a list is passed to a function + # and written into a pure function inside it. + body = Expression( + body.head, body.elements[0], Expression(body.elements[1].head, big) + ) + calls = [0] + original = Expression.replace_slots + + def counting(self, slots, evaluation): + calls[0] += 1 + return original(self, slots, evaluation) + + monkeypatch.setattr(Expression, "replace_slots", counting) + for i in range(5): + new = body.replace_slots([None, session.evaluate(str(i))], session.evaluation) + assert new.elements[1].elements[0] is big + # Three calls per application (the body, #1 and Length[big]), + # not one for each of the 900 subexpressions of big. + assert calls[0] <= 5 * 4, f"{calls[0]} calls for 5 applications" + + +@pytest.mark.parametrize( + ("str_expr", "str_expected"), + [ + ( + "rsData = Range[5]; rsF[l_] := Map[{#, Length[l]} &, {a, b}]; rsF[rsData]", + "{{a, 5}, {b, 5}}", + ), + ("Map[{#, {1, 2}} &, {a, b}]", "{{a, {1, 2}}, {b, {1, 2}}}"), + ("Apply[{##, {x}} &, {a, b}]", "{a, b, {x}}"), + ("Map[Function[{#, # &}], {1}]", "{{1, #1 &}}"), + ("Map[{#, Hold[#]} &, {a}]", "{{a, Hold[a]}}"), + ("ClearAll[rsData, rsF]", "Null"), + ], +) +def test_pure_functions(str_expr, str_expected): + check_evaluation(str_expr, str_expected) + + +@pytest.mark.parametrize( + ("str_expr", "str_expected"), + [ + # Slots in a compound head, at any depth. + ("Map[#[0] &, {Sin, Cos}]", "{0, 1}"), + ("Map[{a, {#[0]}} &, {Cos}]", "{{a, {1}}}"), + ("Map[f[#][x] &, {1, 2}]", "{f[1][x], f[2][x]}"), + ("Map[g[h[#]][x] &, {1}]", "{g[h[1]][x]}"), + ("Apply[f[##][y] &, {1, 2}]", "f[1, 2][y]"), + ], +) +def test_slots_in_compound_heads(str_expr, str_expected): + check_evaluation(str_expr, str_expected)