1138 plot functions for ast in generation package (#1140)

Improved handling and visualization of ast in memilio-generation

Author: dr2001dlr

pycode/memilio-generation/memilio/generation/__init__.py

@@ -25,3 +25,4 @@
 from .intermediate_representation import IntermediateRepresentation
 from .scanner import Scanner
 from .scanner_config import ScannerConfig
+from .ast import AST

pycode/memilio-generation/memilio/generation/ast.py

@@ -0,0 +1,152 @@
+#############################################################################
+# Copyright (C) 2020-2024 MEmilio
+#
+# Authors: Maximilian Betz, Daniel Richter
+#
+# Contact: Martin J. Kuehn <Martin.Kuehn@DLR.de>
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#############################################################################
+"""
+@file ast.py
+@brief Create the ast and assign ids. Get ids and nodes. 
+"""
+import subprocess
+import tempfile
+import logging
+from clang.cindex import Cursor, TranslationUnit, Index, CompilationDatabase
+from typing import TYPE_CHECKING
+from memilio.generation import utility
+
+
+if TYPE_CHECKING:
+    from memilio.generation import ScannerConfig
+
+from typing_extensions import Self
+
+
+class AST:
+    """! Create the ast and assign ids.
+    Functions for getting nodes and node ids.
+    """
+
+    def __init__(self: Self, conf: "ScannerConfig") -> None:
+        self.config = conf
+        self.cursor_id = -1
+        self.id_to_val = dict()
+        self.val_to_id = dict()
+        self.cursor = None
+        self.translation_unit = self.create_ast()
+
+    def create_ast(self: Self) -> TranslationUnit:
+        """! Create an abstract syntax tree for the main model.cpp file with a corresponding CompilationDatabase.
+        A compile_commands.json is required (automatically generated in the build process).
+        """
+        self.cursor_id = -1
+        self.id_to_val.clear()
+        self.val_to_id.clear()
+
+        idx = Index.create()
+
+        file_args = []
+
+        unwanted_arguments = [
+            '-Wno-unknown-warning', "--driver-mode=g++", "-O3", "-Werror", "-Wshadow"
+        ]
+
+        dirname = utility.try_get_compilation_database_path(
+            self.config.skbuild_path_to_database)
+        compdb = CompilationDatabase.fromDirectory(dirname)
+        commands = compdb.getCompileCommands(self.config.source_file)
+        for command in commands:
+            for argument in command.arguments:
+                if argument not in unwanted_arguments:
+                    file_args.append(argument)
+        file_args = file_args[1:-4]
+
+        clang_cmd = [
+            "clang-14", self.config.source_file,
+            "-std=c++17", '-emit-ast', '-o', '-']
+        clang_cmd.extend(file_args)
+
+        try:
+            clang_cmd_result = subprocess.run(
+                clang_cmd, stdout=subprocess.PIPE)
+            clang_cmd_result.check_returncode()
+        except subprocess.CalledProcessError as e:
+            # Capture standard error and output
+            logging.error(
+                f"Clang failed with return code {e.returncode}. Error: {clang_cmd_result.stderr.decode()}")
+            raise RuntimeError(
+                f"Clang AST generation failed. See error log for details.")
+
+        # Since `clang.Index.read` expects a file path, write generated abstract syntax tree to a
+        # temporary named file. This file will be automatically deleted when closed.
+        with tempfile.NamedTemporaryFile() as ast_file:
+            ast_file.write(clang_cmd_result.stdout)
+            translation_unit = idx.read(ast_file.name)
+
+        self._assing_ast_with_ids(translation_unit.cursor)
+
+        logging.info("AST generation completed successfully.")
+
+        return translation_unit
+
+    def _assing_ast_with_ids(self, cursor: Cursor) -> None:
+        """! Traverse the AST and assign a unique ID to each node during traversal.
+
+        @param cursor: The current node (Cursor) in the AST to traverse.
+        """
+        # assing_ids umschreiben -> mapping
+        self.cursor_id += 1
+        id = self.cursor_id
+        self.id_to_val[id] = cursor
+
+        if cursor.hash in self.val_to_id.keys():
+            self.val_to_id[cursor.hash].append(id)
+        else:
+            self.val_to_id[cursor.hash] = [id]
+
+        logging.info(
+            f"Node {cursor.spelling or cursor.kind} assigned ID {id}")
+
+        for child in cursor.get_children():
+            self._assing_ast_with_ids(child)
+
+    @property
+    def root_cursor(self):
+        return self.translation_unit.cursor
+
+    def get_node_id(self, cursor: Cursor) -> int:
+        """! Returns the id of the current node.
+
+        Extracs the key from the current cursor from the dictonary id_to_val
+
+        @param cursor: The current node of the AST as a cursor object from libclang.
+        """
+        for cursor_id in self.val_to_id[cursor.hash]:
+
+            if self.id_to_val[cursor_id] == cursor:
+
+                return cursor_id
+        raise IndexError(f"Cursor {cursor} is out of bounds.")
+
+    def get_node_by_index(self, index: int) -> Cursor:
+        """! Returns the node at the specified index position.
+
+        @param index: Node_id from the ast.
+        """
+
+        if index < 0 or index >= len(self.id_to_val):
+            raise IndexError(f"Index {index} is out of bounds.")
+        return self.id_to_val[index]

pycode/memilio-generation/memilio/generation/graph_visualization.py

@@ -0,0 +1,164 @@
+#############################################################################
+# Copyright (C) 2020-2024 MEmilio
+#
+# Authors: Maximilian Betz, Daniel Richter
+#
+# Contact: Martin J. Kuehn <Martin.Kuehn@DLR.de>
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#############################################################################
+
+import os
+import logging
+from typing import Callable
+from graphviz import Digraph
+from clang.cindex import Cursor
+from memilio.generation.ast import AST
+
+
+class Visualization:
+    """! Class for plotting the abstract syntax tree in different formats.
+    """
+    @staticmethod
+    def output_ast_terminal(ast: AST, cursor: Cursor) -> None:
+        """! Output the abstract syntax tree to terminal.
+
+        @param ast: ast object from AST class.
+        @param cursor: The current node of the AST as a cursor object from libclang.
+        """
+
+        def terminal_writer(level: int, cursor_label: str) -> None:
+            print(indent(level) + cursor_label)
+
+        _output_cursor_and_children(cursor, ast, terminal_writer)
+
+        logging.info("AST-Terminal written.")
+
+    @staticmethod
+    def output_ast_png(cursor: Cursor, max_depth: int, output_file_name: str = 'ast_graph') -> None:
+        """! Output the abstract syntax tree to a .png. Set the starting node and the max depth.
+
+        To save the abstract syntax tree as an png with a starting node and a depth u cann use the following command
+
+        Example command: aviz.output_ast_png(ast.get_node_by_index(1), 2)
+
+        aviz -> instance of the Visualization class.
+
+        ast -> instance of the AST class.
+
+        .get_node_by_index -> get a specific node by id (use .output_ast_formatted to see node ids)
+
+        The number 2 is a example for the depth the graph will show
+
+        @param cursor: The current node of the AST as a cursor object from libclang.
+        @param max_depth: Maximal depth the graph displays.
+        """
+
+        graph = Digraph(format='png')
+
+        _output_cursor_and_children_graphviz_digraph(
+            cursor, graph, max_depth, 0)
+
+        graph.render(filename=output_file_name, view=False)
+
+        output_path = os.path.abspath(f"{output_file_name}.png")
+        logging.info(f"AST-png written to {output_path}")
+
+    @staticmethod
+    def output_ast_formatted(ast: AST, cursor: Cursor, output_file_name: str = 'ast_formated.txt') -> None:
+        """!Output the abstract syntax tree to a file.
+
+        @param ast: ast object from AST class.
+        @param cursor: The current node of the AST as a cursor object from libclang.
+        """
+
+        with open(output_file_name, 'w') as f:
+            def file_writer(level: int, cursor_label: str) -> None:
+                f.write(indent(level) + cursor_label + newline())
+            _output_cursor_and_children(cursor, ast,  file_writer)
+
+        output_path = os.path.abspath(f"{output_file_name}")
+        logging.info(f"AST-formated written to {output_path}")
+
+
+def indent(level: int) -> str:
+    """! Create an indentation based on the level.
+    """
+    return '│   ' * level + '├── '
+
+
+def newline() -> str:
+    """! Create a new line.
+    """
+    return '\n'
+
+
+def _output_cursor_and_children(cursor: Cursor, ast: AST, writer: Callable[[int, str], None], level: int = 0) -> None:
+    """!Generic function to output the cursor and its children with a specified writer.
+
+    @param cursor: The current node of the AST as a libclang cursor object.
+    @param ast: AST object from the AST class.
+    @param writer: Function that takes `level` and `cursor_label` and handles output.
+    @param level: The current depth in the AST for indentation purposes.
+    """
+
+    cursor_id = ast.get_node_id(cursor)
+
+    cursor_kind = f"<CursorKind.{cursor.kind.name}>"
+    file_path = cursor.location.file.name if cursor.location.file else ""
+
+    if cursor.spelling:
+        cursor_label = (f'ID:{cursor_id} {cursor.spelling} '
+                        f'{cursor_kind}   '
+                        f'{file_path}')
+    else:
+        cursor_label = f'ID:{cursor_id} {cursor_kind} {file_path}'
+
+    writer(level, cursor_label)
+
+    for child in cursor.get_children():
+        _output_cursor_and_children(
+            child, ast, writer, level + 1)
+
+
+def _output_cursor_and_children_graphviz_digraph(cursor: Cursor, graph: Digraph, max_d: int, current_d: int, parent_node: str = None) -> None:
+    """! Output the cursor and its children as a graph using Graphviz.
+
+    @param cursor: The current node of the AST as a Cursor object from libclang.
+    @param graph: Graphviz Digraph object where the nodes and edges will be added.
+    @param max_d: Maximal depth.
+    @param current_d: Current depth.
+    @param parent_node: Name of the parent node in the graph (None for the root node).
+    """
+
+    if current_d > max_d:
+        return
+
+    node_label = f"{cursor.kind.name}{newline()}({cursor.spelling})" if cursor.spelling else cursor.kind.name
+
+    current_node = f"{cursor.kind.name}_{cursor.hash}"
+
+    graph.node(current_node, label=node_label)
+
+    if parent_node:
+        graph.edge(parent_node, current_node)
+
+    if cursor.kind.is_reference():
+        referenced_label = f"ref_to_{cursor.referenced.kind.name}{newline()}({cursor.referenced.spelling})"
+        referenced_node = f"ref_{cursor.referenced.hash}"
+        graph.node(referenced_node, label=referenced_label)
+        graph.edge(current_node, referenced_node)
+
+    for child in cursor.get_children():
+        _output_cursor_and_children_graphviz_digraph(
+            child, graph, max_d, current_d + 1, current_node)