Featured Automata learner: now with actual guards. whoop whoop!

1 files changed, 17 insertions, 0 deletions

diff --git a/lib/functions.py b/lib/functions.py
index 187e6ff..35b04ef 100644
--- a/lib/functions.py
+++ b/lib/functions.py
@@ -466,6 +466,23 @@ class SplitFunction(ModelFunction):
         )
         return hyper
 
+    # SplitFunction only
+    def flatten(self):
+        paths = list()
+        for param_value, subtree in self.child.items():
+            if type(subtree) is SplitFunction:
+                for path, value in subtree.flatten():
+                    path = [(self.param_name, param_value)] + path
+                    paths.append((path, value))
+            elif type(subtree) is StaticFunction:
+                path = [(self.param_name, param_value)]
+                paths.append((path, subtree.value))
+            else:
+                raise RuntimeError(
+                    "flatten is only implemented for RMTs with constant leaves"
+                )
+        return paths
+
     @classmethod
     def from_json(cls, data):
         assert data["type"] == "split"