PTA: Support unconditional param setters (e.g. initialization functions)

2 files changed, 38 insertions, 5 deletions

diff --git a/bin/test_automata.py b/bin/test_automata.py
index fd3f462..e537076 100755
--- a/bin/test_automata.py
+++ b/bin/test_automata.py
@@ -31,6 +31,9 @@ example_json_1 = {
             'duration' : {
                 'static' : 50000,
             },
+            'set_param' : {
+                'txpower' : 10
+            },
         },
         {
             'name' : 'setTxPower',
@@ -120,6 +123,7 @@ class TestPTA(unittest.TestCase):
     def test_from_json_function(self):
         pta = PTA.from_json(example_json_1)
         self.assertEqual(pta.states['TX'].get_energy(1000, {'datarate' : 10, 'txbytes' : 6, 'txpower' : 10 }), 1000 * (100 + 2 * 10))
+        self.assertEqual(pta.transitions[3].get_timeout({'datarate' : 10, 'txbytes' : 6, 'txpower' : 10 }), 500 + 16 * 6)
 
     def test_simulation(self):
         pta = PTA()
@@ -212,6 +216,25 @@ class TestPTA(unittest.TestCase):
             'length' : None
         })
 
+    def test_simulation_set_param(self):
+        pta = PTA(parameters = ['txpower', 'length'])
+        pta.add_state('IDLE', power = 5)
+        pta.add_state('TX', power = 100)
+        pta.add_transition('UNINITIALIZED', 'IDLE', 'init', energy = 500000, duration = 50000, set_param = {'txpower' : 10})
+        trace = [
+            ['init'],
+        ]
+        expected_energy = 500000
+        expected_duration = 50000
+        power, duration, state, parameters = pta.simulate(trace)
+        self.assertEqual(power, expected_energy)
+        self.assertEqual(duration, expected_duration)
+        self.assertEqual(state.name, 'IDLE')
+        self.assertEqual(parameters, {
+            'txpower' : 10,
+            'length' : None
+        })
+
     def test_simulation_arg_function(self):
         pta = PTA(parameters = ['txpower', 'length'])
         pta.add_state('IDLE', power = 5)
diff --git a/lib/automata.py b/lib/automata.py
index 4e76581..5e80ad6 100755
--- a/lib/automata.py
+++ b/lib/automata.py
@@ -17,7 +17,7 @@ class Transition:
             timeout = 0, timeout_function = None,
             is_interrupt = False,
             arguments = [], param_update_function = None,
-            arg_to_param_map = None):
+            arg_to_param_map = None, set_param = None):
         self.name = name
         self.origin = orig_state
         self.destination = dest_state
@@ -31,6 +31,7 @@ class Transition:
         self.arguments = arguments.copy()
         self.param_update_function = param_update_function
         self.arg_to_param_map = arg_to_param_map
+        self.set_param = set_param
 
     def get_duration(self, param_dict = {}, args = []):
         if self.duration_function:
@@ -50,12 +51,14 @@ class Transition:
     def get_params_after_transition(self, param_dict, args = []):
         if self.param_update_function:
             return self.param_update_function(param_dict, args)
+        ret = param_dict.copy()
         if self.arg_to_param_map:
-            ret = param_dict.copy()
             for k, v in self.arg_to_param_map.items():
                 ret[k] = args[v]
-            return ret
-        return param_dict
+        if self.set_param:
+            for k, v in self.set_param.items():
+                ret[k] = v
+        return ret
 
 class State:
     def __init__(self, name, power = 0, power_function = None):
@@ -142,11 +145,18 @@ class PTA:
             duration_function = _json_function_to_analytic_function(transition, 'duration', pta.parameters)
             energy_function = _json_function_to_analytic_function(transition, 'energy', pta.parameters)
             timeout_function = _json_function_to_analytic_function(transition, 'timeout', pta.parameters)
+            arg_to_param_map = None
+            if 'arg_to_param_map' in transition:
+                arg_to_param_map = transition['arg_to_param_map']
             pta.add_transition(transition['origin'], transition['destination'],
                 transition['name'],
                 duration = _json_get_static(transition, 'duration'),
+                duration_function = duration_function,
                 energy = _json_get_static(transition, 'energy'),
-                timeout = _json_get_static(transition, 'timeout')
+                energy_function = energy_function,
+                timeout = _json_get_static(transition, 'timeout'),
+                timeout_function = timeout_function,
+                arg_to_param_map = arg_to_param_map
             )
 
         return pta