Add notion of wanted states and a gc mechanism

1 files changed, 26 insertions, 6 deletions

diff --git a/statemachine.py b/statemachine.py
index 41f4570..8be3db7 100644
--- a/statemachine.py
+++ b/statemachine.py
@@ -15,6 +15,7 @@ class StateMachine:
         Create a new state machine
         """
         self.up = set()
+        self.wanted = set()
         self.deps = []
         self.event_triggers = []
 
@@ -28,7 +29,7 @@ class StateMachine:
                 retval |= self.bring_up(cat.intersect_args(**event.args))
         return retval
 
-    def bring_up(self, cat):
+    def bring_up(self, cat, wanted=True):
         """
         Move states in the given Category `cat` from down to up.
         """
@@ -42,11 +43,11 @@ class StateMachine:
             else:
                 found.append(res)
         if found == None:
-            self.add_hold(cat)
+            self.add_hold(cat, wanted)
             return True
         to_add = self.cat_cross(found)
         for x in to_add:
-            self.add_hold(x)
+            self.add_hold(x, wanted)
         return True
 
     def cat_cross(self, found):
@@ -69,7 +70,7 @@ class StateMachine:
             to_add.add(tup[0])
         return to_add
 
-    def add_hold(self, cat):
+    def add_hold(self, cat, wanted):
         """
         Add a hold to a state. Does not check dependencies.
         """
@@ -81,8 +82,9 @@ class StateMachine:
                 self.up.add(cat)
                 return
         self.up.add(cat)
+        if wanted: self.wanted.add(cat)
 
-    def bring_down(self, cat):
+    def bring_down(self, cat, rec=False):
         """
         Bring a currently "up" state down.
         """
@@ -94,8 +96,26 @@ class StateMachine:
             if None in match: match.remove(None)
             if match != None:
                 for item in match:
-                    self.bring_down(dependent.fill(item.args))
+                    self.bring_down(dependent.fill(item.args), rec=True)
         self.up -= to_drop
+        self.wanted -= to_drop
+        if not rec: self.cleanup_states()
+
+    def cleanup_states(self):
+        """
+        Remove unwanted states
+        """
+        new_up = self.wanted.copy()
+        addition = set([1])
+        while len(addition):
+            addition = set()
+            for s in addition:
+                for x in self.get_applicable_deps(s):
+                    for y in self.up:
+                        if y.subset_of(x):
+                            addition.add(y)
+            new_up |= addition
+        self.up = new_up
 
     def get_satisfied_states(self, dependents, dependencies):
         """