ast_generator/ast_generator.py

@@ -25,6 +25,7 @@ class AstGenerator:
     falls into will be selected.
     """
 
+    ### INITIALIZATION ###
     def __init__(self, settings: dict):
         """
         This class is designed to get  the settings from some wrapper class that
@@ -59,7 +60,7 @@ class AstGenerator:
         self.bool_op_options, self.bool_op_cutoffs, self.bool_op_numline = (
             get_numberlines('expression-weights', ['brackets', 'comparison', 'logical', 'unary'],
                             excluded_values=[[], ['less-than-or-equal', 'greater-than-or-equal', 'less-than',
-                                                       'greater-than'], [], ['noop', 'negation']],
+                                                  'greater-than'], [], ['noop', 'negation']],
                             settings=self.settings))
         self.bool_unary = ['not']
         self.float_op_options, self.float_op_cutoffs, self.float_op_numline = (
@@ -83,49 +84,13 @@ class AstGenerator:
         self.variable_names = var_name_list[0:var_name_len // 2]
         self.routine_names = var_name_list[var_name_len // 2:var_name_len]
 
+    ### GENERATION ###
     def generate_ast(self):
         """
         @brief generates an AST from a grammar
         """
         self.generate_top_level_block()
 
-    def make_element(self, name: str, keys: list[tuple[str, any]]) -> ET.Element:
-        """
-        @brief make an xml element for the ast
-
-        @effects modifies self.current_ast_element
-
-        @param name: the tag for the element
-        @param keys: a list of tuple containing keys for the element
-        """
-        element = build_xml_element(keys, name=name)
-        if self.current_ast_element is not None:
-            self.current_ast_element.append(element)
-        self.current_ast_element = element
-
-        return element
-
-    def make_scoped_element(self, name, keys) -> ET.Element:
-        """
-        @brief make an xml element for the ast with a scope
-
-        @param name: the tag for the element
-        @param keys: a list of tuple containing keys for the element
-        """
-        parent = self.current_ast_element
-        self.push_scope()
-        self.make_element(name, keys)
-        return parent
-
-    def exit_scoped_element(self, parent):
-        """
-        @brief leave the current element and return to parent
-
-        @param parent: the enclosing element to return to
-        """
-        self.pop_scope()
-        self.current_ast_element = parent
-
     def generate_top_level_block(self):
         """
         @brief creates the top-level block containing the whole program
@@ -185,103 +150,21 @@ class AstGenerator:
         self.pop_scope()
         self.current_ast_element = parent
 
-    def generate_loop_condition_check(self, loop_var: Variable):
-        """
-        @brief generates the loop condition check
-
-        Ensures that the loop does not iterate more than max-loop-iterations times
-
-        @param loop_var:
-        @return:
-        """
-        # loop var is always an int
-        assert loop_var.type == GAZ_INT_KEY
-
-        # create a conditional xml tag
-        if_stmt = build_xml_element([], name=GAZ_IF_TAG)
-        self.current_ast_element.append(if_stmt)
-        parent = self.current_ast_element
-        self.current_ast_element = if_stmt
-
-        # add the check 'if loop_var >= self.settings['generation_options']['max-loop-iterations']: break'
-        operation = build_xml_element([("op", ">=")], name=GAZ_OPERATOR_TAG)
-        self.current_ast_element.append(operation)
-        self.current_ast_element = operation
-
-        lhs = build_xml_element([], name=GAZ_LHS_TAG)
-        operation.append(lhs)
-
-        var = build_xml_element([("name", loop_var.name), ("type", loop_var.type)], name=GAZ_VAR_TAG)
-        lhs.append(var)
-
-        rhs = build_xml_element([], name=GAZ_RHS_TAG)
-        operation.append(rhs)
-        rhs.append(
-            self.make_literal(GAZ_INT_KEY, "'" + str(self.settings['generation-options']['max-loop-iterations']) + "'"))
-
-        true_block = build_xml_element([], name=GAZ_BLOCK_TAG)
-        if_stmt.append(true_block)
-        self.current_ast_element = true_block
-        break_stmt = build_xml_element([], name=GAZ_BREAK_TAG)
-        true_block.append(break_stmt)
-
-        # return everything to normalcy
-        self.current_ast_element = parent
-
-    def generate_loop_condition_increment(self, loop_var):
-        assert loop_var.type == GAZ_INT_KEY
-
-        parent = self.current_ast_element
-        assignment = build_xml_element([], name=GAZ_ASSIGNMENT_TAG)
-        self.current_ast_element.append(assignment)
-        self.current_ast_element = assignment
-
-        # append the variable
-        self.current_ast_element.append(loop_var.xml)
-
-        # add the increment 'loop_var += 1'
-        assn_rhs = build_xml_element([], name=GAZ_RHS_TAG)
-        self.current_ast_element.append(assn_rhs)
-        self.current_ast_element = assn_rhs
-
-        operation = build_xml_element([("op", "+")], name=GAZ_OPERATOR_TAG)
-        self.current_ast_element.append(operation)
-        self.current_ast_element = operation
-
-        lhs = build_xml_element([], name=GAZ_LHS_TAG)
-        operation.append(lhs)
-
-        var = build_xml_element([("name", loop_var.name), ("type", loop_var.type)], name=GAZ_VAR_TAG)
-        lhs.append(var)
-
-        rhs = build_xml_element([], name=GAZ_RHS_TAG)
-        operation.append(rhs)
-        rhs.append(self.make_literal(GAZ_INT_KEY, '1'))
-
-        # return everything to normalcy
-        self.current_ast_element = parent
-
     def generate_return(self, return_type=None, return_value=None):
         if return_type is None or return_type == GAZ_VOID_TYPE:
-            self.current_ast_element.append(build_xml_element([], name=GAZ_RETURN_TAG))
+            self.current_ast_element.append(self.make_element(GAZ_RETURN_TAG, []))
             return
         else:
+            keys = [("type", return_type)]
+            element = self.make_element(GAZ_RETURN_TAG, keys)
+            self.current_ast_element.append(element)
+            parent = self.current_ast_element
             if return_value is None:
-                xml_element = build_xml_element([("type", return_type)], name=GAZ_RETURN_TAG)
-                self.current_ast_element.append(xml_element)
-                parent = self.current_ast_element
-                self.current_ast_element = xml_element
                 self.generate_expression(return_type)
-                self.current_ast_element = parent
-                return
             else:
-                xml_element = build_xml_element([("type", return_type)], name=GAZ_RETURN_TAG)
-                self.current_ast_element.append(xml_element)
-                parent = self.current_ast_element
-                self.current_ast_element = xml_element
                 self.current_ast_element.append(self.make_literal(return_type, return_value))
-                self.current_ast_element = parent
+
-                return
+            self.current_ast_element = parent
 
     def generate_routine(self, routine_type=None):
         if routine_type is None:
@@ -783,3 +666,110 @@ class AstGenerator:
             for i in range(len(cutoffs)):
                 if res < cutoffs[i]:
                     return types[i]
+
+    ### LOOP HELPERS ###
+
+    def generate_loop_condition_check(self, loop_var: Variable):
+        """
+        @brief generates the loop condition check
+
+        Ensures that the loop does not iterate more than max-loop-iterations times
+
+        @param loop_var:
+        @return:
+        """
+        # loop var is always an int
+        assert loop_var.type == GAZ_INT_KEY
+
+        # create a conditional xml tag
+        if_stmt = build_xml_element([], name=GAZ_IF_TAG)
+        self.current_ast_element.append(if_stmt)
+        parent = self.current_ast_element
+        self.current_ast_element = if_stmt
+
+        # add the check 'if loop_var >= self.settings['generation_options']['max-loop-iterations']: break'
+        operation = build_xml_element([("op", ">=")], name=GAZ_OPERATOR_TAG)
+        rhs = self._loop_heloper(loop_var, operation)
+        rhs.append(
+            self.make_literal(GAZ_INT_KEY, "'" + str(self.settings['generation-options']['max-loop-iterations']) + "'"))
+
+        true_block = build_xml_element([], name=GAZ_BLOCK_TAG)
+        if_stmt.append(true_block)
+        self.current_ast_element = true_block
+        break_stmt = build_xml_element([], name=GAZ_BREAK_TAG)
+        true_block.append(break_stmt)
+
+        # return everything to normalcy
+        self.current_ast_element = parent
+
+    def _loop_heloper(self, loop_var, operation):
+        self.current_ast_element.append(operation)
+        self.current_ast_element = operation
+        lhs = build_xml_element([], name=GAZ_LHS_TAG)
+        operation.append(lhs)
+        var = build_xml_element([("name", loop_var.name), ("type", loop_var.type)], name=GAZ_VAR_TAG)
+        lhs.append(var)
+        rhs = build_xml_element([], name=GAZ_RHS_TAG)
+        operation.append(rhs)
+        return rhs
+
+    def generate_loop_condition_increment(self, loop_var):
+        assert loop_var.type == GAZ_INT_KEY
+
+        parent = self.current_ast_element
+        assignment = build_xml_element([], name=GAZ_ASSIGNMENT_TAG)
+        self.current_ast_element.append(assignment)
+        self.current_ast_element = assignment
+
+        # append the variable
+        self.current_ast_element.append(loop_var.xml)
+
+        # add the increment 'loop_var += 1'
+        assn_rhs = build_xml_element([], name=GAZ_RHS_TAG)
+        self.current_ast_element.append(assn_rhs)
+        self.current_ast_element = assn_rhs
+
+        operation = build_xml_element([("op", "+")], name=GAZ_OPERATOR_TAG)
+        rhs = self._loop_heloper(loop_var, operation)
+        rhs.append(self.make_literal(GAZ_INT_KEY, '1'))
+
+        # return everything to normalcy
+        self.current_ast_element = parent
+
+    ### HELPER FUNCTIONS ###
+    def make_element(self, name: str, keys: list[tuple[str, any]]) -> ET.Element:
+        """
+        @brief make an xml element for the ast
+
+        @effects modifies self.current_ast_element
+
+        @param name: the tag for the element
+        @param keys: a list of tuple containing keys for the element
+        """
+        element = build_xml_element(keys, name=name)
+        if self.current_ast_element is not None:
+            self.current_ast_element.append(element)
+        self.current_ast_element = element
+
+        return element
+
+    def make_scoped_element(self, name, keys) -> ET.Element:
+        """
+        @brief make an xml element for the ast with a scope
+
+        @param name: the tag for the element
+        @param keys: a list of tuple containing keys for the element
+        """
+        parent = self.current_ast_element
+        self.push_scope()
+        self.make_element(name, keys)
+        return parent
+
+    def exit_scoped_element(self, parent):
+        """
+        @brief leave the current element and return to parent
+
+        @param parent: the enclosing element to return to
+        """
+        self.pop_scope()
+        self.current_ast_element = parent