Faster flatten for Sequence symbols

mathics/core/expression.py

@@ -142,6 +142,15 @@ def __new__(cls, *args, **kwargs):
         self.last_evaluated = None
         return self
 
+    def sequences(self):
+        return None
+
+    def flatten_sequence(self):
+        return self
+
+    def flatten_pattern_sequence(self):
+        return self
+
     def get_attributes(self, definitions):
         return set()
 
@@ -498,18 +507,74 @@ def __cmp(self, other):
         return 0
 
 
+def _sequences(leaves):
+    for i, leaf in enumerate(leaves):
+        if leaf.get_head_name() == 'System`Sequence' or leaf.sequences():
+            yield i
+
+
 class Expression(BaseExpression):
     def __new__(cls, head, *leaves):
         self = super(Expression, cls).__new__(cls)
         if isinstance(head, six.string_types):
             head = Symbol(head)
         self.head = head
         self.leaves = [from_python(leaf) for leaf in leaves]
+        self._sequences = None
         return self
 
+    def sequences(self):
+        seq = self._sequences
+        if seq is None:
+            seq = list(_sequences(self.leaves))
+            self._sequences = seq
+        return seq
+
+    def _flatten_sequence(self, sequence):
+        indices = self.sequences()
+        if not indices:
+            return self
+
+        leaves = self.leaves
+
+        flattened = []
+        extend = flattened.extend
+
+        k = 0
+        for i in indices:
+            extend(leaves[k:i])
+            extend(sequence(leaves[i]))
+            k = i + 1
+        extend(leaves[k:])
+
+        return Expression(self.head, *flattened)
+
+    def flatten_sequence(self):
+        def sequence(leaf):
+            if leaf.get_head_name() == 'System`Sequence':
+                return leaf.leaves
+            else:
+                return [leaf]
+
+        return self._flatten_sequence(sequence)
+
+    def flatten_pattern_sequence(self):
+        def sequence(leaf):
+            flattened = leaf.flatten_pattern_sequence()
+            if leaf.get_head_name() == 'System`Sequence' and leaf.pattern_sequence:
+                return flattened.leaves
+            else:
+                return [flattened]
+
+        expr = self._flatten_sequence(sequence)
+        if hasattr(self, 'options'):
+            expr.options = self.options
+        return expr
+
     def copy(self):
         result = Expression(
             self.head.copy(), *[leaf.copy() for leaf in self.leaves])
+        result._sequences = self._sequences
         result.options = self.options
         result.original = self
         # result.last_evaluated = self.last_evaluated
@@ -520,6 +585,7 @@ def shallow_copy(self):
         # the original, only the Expression instance is new.
         expr = Expression(self.head)
         expr.leaves = self.leaves
+        expr._sequences = self._sequences
         expr.options = self.options
         expr.last_evaluated = self.last_evaluated
         return expr
@@ -808,7 +874,7 @@ def eval_range(indices):
 
             if ('System`SequenceHold' not in attributes and    # noqa
                 'System`HoldAllComplete' not in attributes):
-                new = new.flatten(SEQUENCE)
+                new = new.flatten_sequence()
                 leaves = new.leaves
 
             for leaf in leaves:
@@ -1461,9 +1527,6 @@ def __getnewargs__(self):
         return (self.name, self.sympy_dummy)
 
 
-SEQUENCE = Symbol('Sequence')
-
-
 class Number(Atom):
     def __str__(self):
         return str(self.value)

mathics/core/rules.py

@@ -4,7 +4,7 @@
 from __future__ import unicode_literals
 from __future__ import absolute_import
 
-from mathics.core.expression import Expression, SEQUENCE, strip_context, KeyComparable
+from mathics.core.expression import Expression, strip_context, KeyComparable
 from mathics.core.pattern import Pattern, StopGenerator
 
 
@@ -46,17 +46,7 @@ def yield_match(vars, rest):
                 result = new_expression
 
             # Flatten out sequences (important for Rule itself!)
-
-            def flatten(expr):
-                new_expr = expr.flatten(SEQUENCE, pattern_only=True)
-                if not new_expr.is_atom():
-                    for index, leaf in enumerate(new_expr.leaves):
-                        new_expr.leaves[index] = flatten(leaf)
-                if hasattr(expr, 'options'):
-                    new_expr.options = expr.options
-                return new_expr
-
-            result = flatten(result)
+            result = result.flatten_pattern_sequence()
             if return_list:
                 result_list.append(result)
                 # count += 1