modify_alias

#!/usr/bin/env python3
import argparse
from typing import Sequence
from subprocess import Popen, PIPE, run, DEVNULL
import os
import time

import zmq


class InteractiveAliasAnalysisDriver:
    args: argparse.Namespace
    """
    The argument parsers namespace which holds the parsed commandline
    attributes.
    """

    def register_all_arguments(self,
                               arg_parser: argparse.ArgumentParser) -> None:
        """
        Registers all the command line arguments that are used by this tool.

        Add other/additional arguments by overloading this function.
        """
        arg_parser.add_argument("input_file",
                                type=str,
                                nargs="?",
                                help="path to input file")

        arg_parser.add_argument(
            "-op",
            "--optpath",
            type=str,
            help="full path to the opt binary",
            default="/home/michel/ETH/AST/llvm-project/build_interact/bin/opt")

        arg_parser.add_argument(
            "-cp",
            "--compilerpath",
            type=str,
            help=
            "full path to the compiler binary (needs to understand LLVM IR)",
            default="/home/michel/ETH/AST/llvm-project/build_ast/bin/clang")

        arg_parser.add_argument("-csmith",
                                help="Use if the file is generated by csmith",
                                action='store_true')

    def __init__(self, args: Sequence[str] | None = None):

        # arg handling
        arg_parser = argparse.ArgumentParser()
        self.register_all_arguments(arg_parser)
        self.args = arg_parser.parse_args(args=args)

        context = zmq.Context()
        self.socket = context.socket(zmq.REP)
        self.socket.bind("ipc:///tmp/0")

    def run(self) -> None:
        """
        Runs the tool.
        """
        #self.measure_all_modification_style(substitutions={
        #    0: 0,
        #    1: 0,
        #    2: 0,
        #    3: 0,
        #})
        #self.replace_single_query(3, 3)

        substitutions = {
            0: 3,
            1: 3,
            2: 3,
            3: 3,
        }

        self.measure_spec_with_substitutions(substitutions, "build_python")

    def get_env(self) -> dict[str, str]:
        """
        Returns the environment variables that are used by the tool.
        """
        my_env = os.environ.copy()
        my_env["CC"] = self.args.compilerpath
        my_env["CXX"] = self.args.compilerpath + "++"
        return my_env

    def setup_spec(self, folder_path: str, my_env: dict[str, str]):

        cmd = "cmake .. -G Ninja -DCMAKE_CXX_FLAGS=\"-Os\" -DCMAKE_C_FLAGS=\"-Os\" -DCMAKE_C_COMPILER=" + self.args.compilerpath + " -DCMAKE_CXX_COMPILER=" + self.args.compilerpath + "++"
        #cmd = "./spec build build_python 600 /home/michel/ETH/AST/llvm-project/build_interact/bin/clang -c"
        print(cmd)
        run(("mkdir " + folder_path).split(" "), cwd="../specbuilder/")

        p = Popen(cmd.split(" "),
                  cwd="../specbuilder/" + folder_path,
                  env=my_env)

        while p.poll() is None:

            try:
                #  Wait for next request from client
                message = list(self.socket.recv(flags=zmq.NOBLOCK))[0]
            except zmq.error.Again:
                continue

            #  Send reply back to client
            self.socket.send_string(str(message))

    def compile_spec_with_substitutions(self, substitutions: dict[int, int],
                                        my_env: dict[str, str],
                                        folder_path) -> None:
        """
        Compiles the spec benchmark.
        """

        p = Popen("ninja -j1 605".split(" "),
                  cwd="../specbuilder/" + folder_path,
                  env=my_env)

        index = 0
        while p.poll() is None:

            try:
                #  Wait for next request from client
                message = list(self.socket.recv(flags=zmq.NOBLOCK))[0]
            except zmq.error.Again:
                continue

            index = index + 1
            #if index % 100000 == 0:
            #    print("Modifying " + str(index))

            #  Send reply back to client
            self.socket.send_string(str(substitutions[message]))
        print("index is " + str(index))

    def measure_spec_with_substitutions(self, substitutions: dict[int, int],
                                        folder_path: str):

        my_env = self.get_env()

        self.setup_spec(folder_path, my_env)

        self.compile_spec_with_substitutions(substitutions, my_env,
                                             folder_path)

        run([
            "python3", "compare_sizes.py", "build_with_exceptions/",
            folder_path + "/", "605"
        ],
            cwd="../specbuilder")

    # Replaces to_replace with replacement
    def replace_single_query(self, to_replace: int, replacement: int) -> bool:
        self.measure_single_modification_style(
            to_replace,
            modification_callback=lambda cmd, index: self.
            execute_with_modifications(cmd, index, replacement, to_replace))

    # Measures the map all occurences style.
    def measure_all_modification_style(self, substitutions: dict[str, int]):
        self.get_ll_file(self.args.compilerpath, self.args.input_file)

        cmd = [self.args.optpath, '-Os', 'csmith/file2.ll', '-S']

        count = self.get_count_of_queries(cmd, 3)

        print("count is " + str(count))

        self.compile_with_substitutions(cmd, substitutions,
                                        "files_sache/file0.ll")

        print("Done with modifications")

        # Compile groundtruth
        compile_cmd = [
            self.args.compilerpath, 'files_sache/truth.ll', '-o',
            "files_sache/truth.out"
        ]
        print(compile_cmd)
        run(compile_cmd, stdout=DEVNULL, stderr=DEVNULL, text=True)

        if not self.compile_file(0):
            print("Compilation failed for modified")

        print("Done with compilation")

        # Measure groundtruth
        true_size = self.measure_outputsize("files_sache/truth.out")
        print("True size is " + str(true_size))

        size = self.measure_outputsize("files_sache/file0.out")
        if size == -1:
            print("Size measurement failed for modified")
        print("new size: " + str(size))

    # Measures the nth occurence replacement style.
    def measure_single_modification_style(self, to_replace: int,
                                          modification_callback):
        self.get_ll_file(self.args.compilerpath, self.args.input_file)

        cmd = [self.args.optpath, '-Os', 'csmith/file2.ll', '-S']

        count = self.get_count_of_queries(cmd, to_replace)

        print("Count is " + str(count))

        for i in range(count):
            if i % 10 == 0:
                print("Modifying " + str(i))
            if not modification_callback(cmd, i):
                print("Modification failed for " + str(i))

        print("Done with modifications")

        # Compile groundtruth
        run([
            self.args.compilerpath, 'files_sache/truth.ll', '-o',
            "files_sache/truth.out"
        ],
            stdout=DEVNULL,
            stderr=DEVNULL,
            text=True)

        for i in range(count):
            if i % 10 == 0:
                print("Compiling " + str(i))
            if not self.compile_file(i):
                print("Compilation failed for " + str(i))

        print("Done with compilation")

        # Measure groundtruth
        true_size = self.measure_outputsize("files_sache/truth.out")
        print("True size is " + str(true_size))

        for i in range(count):
            size = self.measure_outputsize("files_sache/file" + str(i) +
                                           ".out")
            if size == -1:
                print("Size measurement failed for " + str(i))
            elif size != true_size:
                print(str(i) + ": " + str(size))

    # Runs the given compile command mapping AliasResults through substitutions
    #
    # Returns true if the modification was successful
    def compile_with_substitutions(self, cmd: list[str],
                                   substitutions: dict[int, int],
                                   output_file_name: str) -> bool:
        print(cmd + [" substitutions!"])
        p = Popen(cmd, stdin=PIPE, stdout=PIPE, text=True)

        while p.poll() is None:

            try:
                #  Wait for next request from client
                message = list(self.socket.recv(flags=zmq.NOBLOCK))[0]
            except zmq.error.Again:
                continue

            #  Send reply back to client
            self.socket.send_string(str(substitutions[message]))

        output_file = open(output_file_name, 'w')
        output_file.write(p.stdout.read())

        return True

    # Uses the given compiler get the .ll file in csmith/file2.ll for the given
    # input file
    def get_ll_file(self, compiler_path: str, input_file: str) -> None:
        cmd = compiler_path + " -S -emit-llvm -Xclang -disable-llvm-passes -o csmith/file2.ll " + input_file + " -O1" + (
            " -I/usr/include/csmith-2.3.0" if self.args.csmith else "")
        print(cmd)
        run(cmd.split(" "), stdout=DEVNULL, stderr=DEVNULL, text=True)

    # Given a command that runs opt, returns the number of queries with
    # to_replace as result
    #
    # Returns -1 if the command or an input fails Write result into
    # `files_sache/truth.ll`
    def get_count_of_queries(self, cmd: list[str], to_replace: int) -> int:

        print(cmd + [" counting!"])
        p = Popen(cmd, stdout=PIPE, text=True)
        count = 0

        while p.poll() is None:

            try:
                #  Wait for next request from client
                message = list(self.socket.recv(flags=zmq.NOBLOCK))[0]
            except zmq.error.Again:
                continue

            if message == to_replace:
                count = count + 1

            #  Send reply back to client
            self.socket.send_string(str(message))

        output_file = open("files_sache/truth.ll", 'w')
        output_file.write(p.stdout.read())

        return count

    # Changes the nth occurence of to_replace in the given command to change_to
    #
    # returns true if the modification was successful
    def replace_nth_occurence(self, cmd: str, to_replace: int,
                              modify_index: int, change_to: int) -> bool:

        print(cmd + [" replacing occurences!"])
        p = Popen(cmd, stdin=PIPE, stdout=PIPE, text=True)
        index = 0

        while p.poll() is None:
            try:
                #  Wait for next request from client
                message = list(self.socket.recv(flags=zmq.NOBLOCK))[0]
            except zmq.error.Again:
                continue

            #  Send reply back to client
            if message == to_replace and index == modify_index:
                self.socket.send_string(str(change_to))
            else:
                self.socket.send_string(str(message))

            if message == to_replace:
                index = index + 1

        output_file = open("files_sache/file" + str(modify_index) + ".ll", 'w')
        output_file.write(p.stdout.read())

        return True

    # Changes the nth occurence of MayAlias in the given command to change_to
    #
    # Returns true if the modification was successful
    #
    # change_to is 0 for NoAlias, 1 for MusAlias, 2 for PartialAlias, 3 for
    # MayAlias
    def execute_with_modifications(self, cmd: list[str], modify_index: int,
                                   change_to: int, to_replace: int) -> bool:
        return self.replace_nth_occurence(cmd, to_replace, modify_index,
                                          change_to)

    def compile_file(self, index: int) -> bool:
        cmd = [
            self.args.compilerpath, "files_sache/file" + str(index) + ".ll",
            '-o', "files_sache/file" + str(index) + ".out"
        ]
        print(cmd)
        p = run(cmd, stdout=DEVNULL, stderr=DEVNULL, text=True)
        return p.returncode == 0

    def measure_outputsize(self, file: str) -> int:
        cmd = ['llvm-size', file]
        print(cmd)
        p = run(cmd, stdout=PIPE, stderr=PIPE, text=True)
        if p.stderr != "":
            print(p.stderr)
            return -1
        second_line = p.stdout.split("\n")[1]
        line_list = second_line.split("\t")
        # The other results are more info on where size is: test, data, bss
        return int(line_list[3])

    # Changes the first alias queries to alias_results
    #
    # Afterwards, all other alias queries are changed to default
    #
    # if default is -1, the query stays the same
    def compile_with_list(self,
                          cmd: list[str],
                          alias_results: list[int],
                          output_file_name: str,
                          default: int = 3) -> bool:
        print(cmd + [" list-based!"])
        p = Popen(cmd, stdout=PIPE, text=True)
        index = 0

        while p.poll() is None:

            try:
                #  Wait for next request from client
                message = list(self.socket.recv(flags=zmq.NOBLOCK))[0]
            except zmq.error.Again:
                continue

            #  Send reply back to client
            if index < len(alias_results):
                self.socket.send_string(str(alias_results[index]))
            elif default == -1:
                self.socket.send_string(str(message))
            else:
                self.socket.send_string(str(default))
            index = index + 1

        output_file = open(output_file_name, 'w')
        output_file.write(p.stdout.read())

        return True


if __name__ == "__main__":
    InteractiveAliasAnalysisDriver().run()