From 46c1c0f8fed2f58ada127df634922aa246752979 Mon Sep 17 00:00:00 2001 From: bugariua Date: Wed, 24 Mar 2021 11:43:13 +0100 Subject: [PATCH] cached all_args --- pysmt/fnode.py | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/pysmt/fnode.py b/pysmt/fnode.py index ed019c284..715433791 100644 --- a/pysmt/fnode.py +++ b/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]