cached all_args

pysmt/fnode.py

@@ -17,7 +17,9 @@
 #
 """FNode are the building blocks of formulae."""
 import collections
-from typing import List
+from typing import List, Set
+
+from sortedcontainers import SortedSet
 
 import pysmt
 import pysmt.smtlib
@@ -74,7 +76,8 @@ class FNode(object):
     FormulaManager it belongs.
     """
     __slots__ = ["_content", "_node_id",
-                 "_cached_representation", "_is_fun_application", "_is_symbol", "_function_name", "_original"]
+                 "_cached_representation", "_is_fun_application", "_is_symbol",
+                 "_function_name", "_original", "_all_args"]
 
     def __init__(self, content, node_id):
         self._content = content
@@ -86,6 +89,7 @@ def __init__(self, content, node_id):
         if self._is_fun_application:
             self._function_name = self._content.payload._content.payload[0]
         self._original = None  # original none, before simplification
+        self._all_args = None  # all the arguments, computed recursively
         return
 
     # __eq__ is left as default while __hash__ uses the node id. This
@@ -106,6 +110,18 @@ def args(self):
         """Returns the subformulae."""
         return self._content.args
 
+    def all_args(self):
+        """Returns all the subformulae (i.e., also the args of the args, etc.)."""
+        if self._all_args is None:
+            self._all_args: Set['FNode'] = SortedSet([], key=lambda x: True)
+            self._collect_in_set(self, self._all_args)
+        return self._all_args
+
+    def _collect_in_set(self, node: 'FNode', result: Set['FNode']):
+        result.add(node)
+        for child in node.args():
+            self._collect_in_set(child, result)
+
     def arg(self, idx):
         """Return the given subformula at the given position."""
         return self._content.args[idx]