Add basic auto-generated energy accounting code for online model eval

2 files changed, 130 insertions, 11 deletions

diff --git a/lib/automata.py b/lib/automata.py
index d26c0bc..4cd694d 100755
--- a/lib/automata.py
+++ b/lib/automata.py
@@ -522,6 +522,14 @@ class PTA:
         """Return PTA-specific ID of transition."""
         return self.transitions.index(transition)
 
+    def get_state_names(self):
+        """Return lexically sorted list of PTA state names."""
+        return sorted(self.state.keys())
+
+    def get_state_id(self, state: State) -> int:
+        """Return PTA-specific ID of state."""
+        return self.get_state_names().index(state.name)
+
     def get_initial_param_dict(self):
         return dict([[self.parameters[i], self.initial_param_values[i]] for i in range(len(self.parameters))])
 
diff --git a/lib/codegen.py b/lib/codegen.py
index b383f20..fd1b24e 100644
--- a/lib/codegen.py
+++ b/lib/codegen.py
@@ -1,17 +1,24 @@
 """Code generators for multipass dummy drivers for online model evaluation."""
 
+from automata import PTA
+
 header_template = """
 #ifndef DFATOOL_{name}_H
 #define DFATOOL_{name}_H
 
+{includes}
+
 class {name}
 {{
 private:
 {name}(const {name} &copy);
+{private_variables}
+{private_functions}
 
 public:
 {enums}
-{functions}
+{public_variables}
+{public_functions}
 }};
 
 extern {name} {name_lower};
@@ -27,10 +34,98 @@ implementation_template = """
 {name} {name_lower};
 """
 
+array_template = """
+{type} const {name}[{length}] = {{{elements}}};
+"""
+
+class ClassFunction:
+    def __init__(self, class_name, return_type, name, arguments, body):
+        """
+        Create a new C++ class method wrapper.
+
+        :param class_name: Class name
+        :param return_type: function return type
+        :param name: function name
+        :param arguments: list of arguments (must contain type and name)
+        :param body: function body (str)
+        """
+        self.class_name = class_name
+        self.return_type = return_type
+        self.name = name
+        self.arguments = arguments
+        self.body = body
+
+    def get_definition(self):
+        return '{} {}({});'.format(self.return_type, self.name, ', '.join(self.arguments))
+
+    def get_implementation(self):
+        if self.body is None:
+            return ''
+        return '{} {}::{}({}) {{\n{}}}\n'.format(self.return_type, self.class_name, self.name, ', '.join(self.arguments), self.body)
+
+class AccountingMethod:
+    def __init__(self, class_name: str, pta: PTA, ):
+        self.class_name = class_name
+        self.pta = pta
+        self.include_paths = list()
+        self.private_variables = list()
+        self.public_variables = list()
+        self.private_functions = list()
+        self.public_functions = list()
+    
+    def pre_transition_hook(self, transition):
+        return ''
+
+    def init_code(self):
+        return ''
+
+    def get_includes(self):
+        return map(lambda x: '#include "{}"'.format(x), self.include_paths)
+
+class StaticStateOnlyAccounting(AccountingMethod):
+    def __init__(self, class_name: str, pta: PTA, ts_type = 'unsigned int', power_type = 'unsigned int', energy_type = 'unsigned long'):
+        super().__init__(class_name, pta)
+        self.ts_type = ts_type
+        self.include_paths.append('driver/uptime.h')
+        self.private_variables.append('unsigned char lastState;')
+        self.private_variables.append('{} lastStateChange;'.format(ts_type))
+        self.private_variables.append(array_template.format(
+            type = power_type,
+            name = 'state_power',
+            length = len(pta.state),
+            elements = ', '.join(map(lambda state_name: str(pta.state[state_name].power), pta.get_state_names()))
+        ))
+        self.private_variables.append('{} timeInState[{}];'.format(ts_type, len(pta.state)))
+
+        get_energy_function = """
+        {energy_type} total_energy = 0;
+        for (int i = 0; i < {num_states}; i++) {{
+            total_energy += timeInState[i] * state_power[i];
+        }}
+        return total_energy;
+        """.format(energy_type = energy_type, num_states = len(pta.state))
+        self.public_functions.append(ClassFunction(class_name, energy_type, 'getEnergy', list(), get_energy_function))
+
+    def pre_transition_hook(self, transition):
+        return """
+        unsigned int now = uptime.get_us();
+        timeInState[lastState] += now - lastStateChange;
+        lastStateChange = now;
+        lastState = {};
+        """.format(self.pta.get_state_id(transition.destination))
+
+    def init_code(self):
+        return """
+        for (unsigned char i = 0; i < {num_states}; i++) {{
+            timeInState[i] = 0;
+        }}
+        """.format(num_states = len(self.pta.state))
+
+
 class MultipassDriver:
     """Generate C++ header and no-op implementation for a multipass driver based on a DFA model."""
 
-    def __init__(self, name, pta, class_info, enum = dict()):
+    def __init__(self, name, pta, class_info, enum = dict(), accounting = AccountingMethod):
         self.impl = ''
         self.header = ''
         self.name = name
@@ -38,10 +133,13 @@ class MultipassDriver:
         self.class_info = class_info
         self.enum = enum
 
-        function_definitions = list()
-        function_bodies = list()
+        includes = list()
+        private_functions = list()
+        public_functions = list()
+        private_variables = list()
+        public_variables = list()
 
-        function_definitions.append('{}() {{}}'.format(self.name))
+        public_functions.append(ClassFunction(self.name, '', self.name, list(), accounting.init_code()))
         seen_transitions = set()
 
         for transition in self.pta.transitions:
@@ -64,17 +162,30 @@ class MultipassDriver:
             function_definition = '{} {}({})'.format(function_info.return_type, transition.name, ', '.join(function_arguments))
             function_head = '{} {}::{}({})'.format(function_info.return_type, self.name, transition.name, ', '.join(function_arguments))
 
-            function_body = str()
+            function_body = accounting.pre_transition_hook(transition)
 
             if function_info.return_type != 'void':
-                function_body = 'return 0;\n'
+                function_body += 'return 0;\n'
 
-            function_definitions.append(function_definition + ';')
-            function_bodies.append('{} {{\n{}}}'.format(function_head, function_body))
+            public_functions.append(ClassFunction(self.name, function_info.return_type, transition.name, function_arguments, function_body))
 
         enums = list()
         for enum_name in self.enum.keys():
             enums.append('enum {} {{ {} }};'.format(enum_name, ', '.join(self.enum[enum_name])))
 
-        self.header = header_template.format(name = self.name, name_lower = self.name.lower(), functions = '\n'.join(function_definitions), enums = '\n'.join(enums))
-        self.impl = implementation_template.format(name = self.name, name_lower = self.name.lower(), functions = '\n\n'.join(function_bodies))
+        if accounting:
+            includes.extend(accounting.get_includes())
+            private_functions.extend(accounting.private_functions)
+            public_functions.extend(accounting.public_functions)
+            private_variables.extend(accounting.private_variables)
+            public_variables.extend(accounting.public_variables)
+
+        self.header = header_template.format(
+            name = self.name, name_lower = self.name.lower(),
+            includes = '\n'.join(includes),
+            private_variables = '\n'.join(private_variables),
+            public_variables = '\n'.join(public_variables),
+            public_functions = '\n'.join(map(lambda x: x.get_definition(), public_functions)),
+            private_functions = '',
+            enums = '\n'.join(enums))
+        self.impl = implementation_template.format(name = self.name, name_lower = self.name.lower(), functions = '\n\n'.join(map(lambda x: x.get_implementation(), public_functions)))