<gh_stars>1-10
from pysv.smtlib.synth import SynthesisConstr
from pysv.smtlib.synth import SynthesisConstrTestCases
from pysv.templates import *
from pysv.contract import *
from pysv import solvers
from pysv import smt_common



#
# Python ast package documentation: https://docs.python.org/2/library/ast.html
# Python _ast package documentation: not found
# Some useful information: https://julien.danjou.info/blog/2015/python-ast-checking-method-declaration
#
# To run synthesis functions required is Z3 solver (installation instruction in README.txt). Some links:
# - Z3 project main page: https://github.com/Z3Prover/z3
# - general Z3 tutorial: http://rise4fun.com/z3/tutorial/guide
# - python Z3 tutorial: http://cpl0.net/~argp/papers/z3py-guide.pdf
# - Advanced python materials: http://www.cs.tau.ac.il/~msagiv/courses/asv/z3py/
#
# *** z3 binding for Scala: http://lara.epfl.ch/~psuter/ScalaZ3/
#
# z3.parse_smt2_string() - a function which can use smt2 code instead of explicitly using python api.
# (set-option :macro-finder true)  - allows for filling bodies of functions.
#


class SynthesizerSMT(object):

    def __init__(self, code, code_pre, code_post, program_vars, env, holes_defs, assertions = None):
        assert (type(holes_defs) == list), 'Holes definitions should be contained in the list!'
        if assertions is None:
            assertions = []

        self.code = code
        self.code_pre = code_pre
        self.code_post = code_post
        self.program_vars = program_vars
        self.holes_defs = holes_defs
        self.env = env
        self.assertions = assertions
        self.exclude_assertions = []
        self.result = None

    def synthesize(self, assertions = None):
        """Synthesizes a program from specification using only user-defined assertions."""
        if assertions is None:
            assertions = self.assertions
        self.result = synthesize(self.code, self.code_pre, self.code_post, self.program_vars,
                                 self.env, self.holes_defs, assertions)
        return self.result

    def next_program(self):
        """Synthesizes a new program different from all programs previously generated by this method. This is done by adding assertions for excluding combinations of productions used in previous models."""
        if self.result is not None:
            a = SynthesizerSMT.get_assertion_excluding_model(self.result)
            self.exclude_assertions.append(a)
        return self.synthesize(self.all_assertions())

    def find_all_possible_programs(self):
        """Returns a list of all programs meeting the specification."""
        programs = []
        while True:
            res = self.next_program()
            if res.decision != 'sat':
                break
            programs.append(res.final_code)
        return programs

    def reset(self):
        """Clears all result-oriented data produced in this instance since initialization."""
        self.exclude_assertions = []
        self.result = None

    def all_assertions(self):
        """Returns list with both user-defined assertions and internal assertions added during searching for the next program."""
        res = []
        res.extend(self.assertions)
        res.extend(self.exclude_assertions)
        return res

    @staticmethod
    def get_assertion_excluding_model(result):
        used_vars = result.used_struct_vars
        text = '(assert (or '
        for v in used_vars:
            text += '(not (= ' + v + ' ' + result[v] + ')) '
        text += '))'
        return text




def _synthesize_universal(smtlib_constr, code, code_pre, code_post, program_vars, env, holes_decls = None, free_vars = None):
    ib_smt2, pre_smt2, post_smt2 = smt_common.get_code_in_smt2(code, code_pre, code_post, program_vars, env,
                                                               holes_decls)
    script = smtlib_constr.produce_script_synthesis(ib_smt2, pre_smt2, post_smt2, program_vars,
                                                    holes_decls, free_vars=free_vars)
    smt_common.write_script_to_file(script, env)
    if env.only_script:
        print(script)
        return None
    else:
        utils.logger.info('\n\n******** SCRIPT ********:\n' + script)
        # Solving constraints and interpreting result.
        res = solvers.run_solver(script, env)
        synth_res = SynthesisResult(res, code, holes_decls, env)
        return synth_res




def synthesize(code, code_pre, code_post, program_vars, env, holes_decls = None, assertions = None, free_vars = None):
    """This synthesizer executes in an external solver SMT-LIB 2.0 script containing
    generated constraints. Program is synthesized to be correct with respect to delivered pre- and
    post-conditions.

    :param code: (str) Source code of the Python program.
    :param code_pre: (str) Source code of the Python expression representing all *pre-conditions*.
    :param code_post: (str) Source code of the Python expression representing all *post-conditions*.
    :param program_vars: (ProgramVars) Object containing information about variables and their types.
    :param env: (Options) Synthesis options.
    :param holes_decls: (list[HoleDecl]) List containing declarations of all holes present in the program. If empty, then synthesis is performed only on the free variables.
    :param assertions: (list[str]) Optional list of SMT-LIB 2.0 commands, which will be appended at the end of the script.
    :param free_vars: (list[str]) Names of variables in the program which are free and are part of the final solution together with decisions of how holes should be filled.
    :return: (SynthesisResult) Interpreted result of the solver.
    """
    if holes_decls is None:
        holes_decls = []
    if assertions is None:
        assertions = env.assertions
    assert (type(holes_decls) == list), 'Holes definitions should be contained in the list!'

    smtlib = SynthesisConstr(env, assertions)
    return _synthesize_universal(smtlib, code, code_pre, code_post, program_vars, env, holes_decls, free_vars)




def synthesize_tc(test_cases, code, code_pre, code_post, program_vars, env, holes_decls = None, assertions = None, free_vars = None):
    """This synthesizer executes in an external solver SMT-LIB 2.0 script containing
    generated constraints. Program is synthesized to be correct with respect to delivered pre- and
    post-conditions.

    :param test_cases: (TestCases) Set of test cases.
    :param code: (str) Source code of the Python program.
    :param code_pre: (str) Source code of the Python expression representing all *pre-conditions*.
    :param code_post: (str) Source code of the Python expression representing all *post-conditions*.
    :param program_vars: (ProgramVars) Object containing information about variables and their types.
    :param env: (Options) Synthesis options.
    :param holes_decls: (list[HoleDecl]) List containing declarations of all holes present in the program. If empty, then synthesis is performed only on the free variables.
    :param assertions: (list[str]) Optional list of SMT-LIB 2.0 commands, which will be appended at the end of the script.
    :param free_vars: (list[str]) Names of variables in the program which are free and are part of the final solution together with decisions of how holes should be filled.
    :return: (SynthesisResult) Interpreted result of the solver.
    """
    if holes_decls is None:
        holes_decls = []
    if assertions is None:
        assertions = env.assertions
    assert (type(holes_decls) == list), 'Holes definitions should be contained in the list!'

    smtlib = SynthesisConstrTestCases(test_cases, env, assertions, holes_decls)
    return _synthesize_universal(smtlib, code, code_pre, code_post, program_vars, env, holes_decls, free_vars)







class HoleDecl(object):
    """HoleDecl contains all important information about a single hole in the program.

    Attributes:
    -----------
    id : string
        Unique name of the hole.
    grammar : Grammar
        Grammar defining possible hole's content.
    program_vars : ProgramVars | dict[string,string]
        Information about variables present in the grammar and their types. If not provided, variable names and types will be automatically inferred from grammar, but this is not possible in some cases (e.g. when variable is part of rhs of a production but is not a terminal in this production: 'P ::= (+ x P)').
    vars : dict[string,string]
        Dictionary of all variables, input or local, which may be used in this hole.
    is_expression_hole : bool
        Specifies, if this is an expression-hole (e.g. condition in the IF) rather than instruction-hole.
    max_depth : int
        Maximum depth of expressions produced from grammar.
    expr_type : string
        Returned type of this hole, which is extracted from the grammar.
    """
    DEFAULT_MAX_DEPTH = 2
    def __init__(self, id, grammar, program_vars = None, is_expression_hole = True, max_depth = DEFAULT_MAX_DEPTH):
        self.id = id
        self.grammar = grammar
        self.is_expression_hole = is_expression_hole
        self.expr_type = self.grammar['Start'].sort
        self.max_depth = max_depth
        self.start_function_name = HoleGrammarTree.base_function_name(self) + 'Start0'
        if program_vars is None:
            program_vars = ProgramVars(self.grammar.get_used_var_names())
        elif type(program_vars) == dict:
            program_vars = ProgramVars(program_vars)
        assert type(program_vars) == ProgramVars, 'Incorrect variables description passed to HoleDecl!'
        self.program_vars = program_vars
        self.vars = self.program_vars.all()
        HoleDecl.update_grammar_with_variables(self.grammar, self.program_vars)
        HoleDecl.fix_grammar_variables_tags(self.grammar)
        self.grammar_tree = HoleGrammarTree(self, self.max_depth)

    def rename_var(self, old_id, new_id):
        self.program_vars.rename_var(old_id, new_id)
        self.vars = self.program_vars.all()

    def get_used_var_names(self):
        """Returns dictionary containing variable names used in the grammar."""
        # return self.grammar.get_used_var_names()
        return self.vars

    def get_function_call(self):
        """Returns SMT-LIB 2.0 code to insert in place of this hole's symbol."""
        var_names = sorted(self.get_used_var_names().keys())
        return '(' + self.start_function_name + ' ' + ' '.join(var_names) + ')'

    @staticmethod
    def update_grammar_with_variables(grammar, program_vars):
        """Adds to the grammar additional rules for variables of certain types. This is needed for '*Variable' tokens in the SYGUS grammar format."""
        def update(vars, rule_name, tpe):
            filtered_vars = ProgramVars.get_vars_of_type(vars, tpe)
            gr = GrammarRule(rule_name, tpe, grammar)
            for v in sorted(filtered_vars.keys()):
                p = Production(rule_name, [v], gr)
                gr.add(p)
            grammar.add_rule(gr)

        def update_all(vars, rule_name):
            for tpe in ProgramVars.get_types(vars):
                update(vars, rule_name+'_'+str(tpe), tpe)

        update_all(program_vars.input_vars, 'InputVariable')
        update_all(program_vars.local_vars, 'LocalVariable')
        update_all(program_vars.all(), 'Variable')

    @staticmethod
    def fix_grammar_variables_tags(grammar):
        """Replaces certain tags for variables ('InputVariable', 'LocalVariable', 'Variable') with names of their appropriate grammar symbols. Works in place."""
        def search(rule, name):
            for p in rule:
                if name in p.body:
                    i = p.body.index(name)
                    vtype = p.body[i+1]
                    c = utils.index_of_closing_parenthesis(p.body, i-1)
                    del p.body[i-1 : c+1]
                    p.body.insert(i-1, name + '_' + vtype)

        for r in grammar:
            search(grammar[r], 'InputVariable')
            search(grammar[r], 'LocalVariable')
            search(grammar[r], 'Variable')

    @classmethod
    def from_str(cls, text):
        """Creates a hole declaration from a string encoding. Example of the string encoding: "H0,((Start Bool (a b)))"."""
        w = text.split(",")
        assert len(w) <= 3 , "Too many fields in the description of a hole."
        assert len(w) >= 2 , "Hole must be described with at least two fields."
        id = w[0].strip()
        if len(w) == 2:
            depth = HoleDecl.DEFAULT_MAX_DEPTH
            grammar = load_gramar_from_SYGUS_spec(w[1])
        else: # len(w) == 3
            depth = int(w[1])
            grammar = load_gramar_from_SYGUS_spec(w[2])
        return HoleDecl(id, grammar, max_depth=depth)

    @classmethod
    def many_from_str(cls, text):
        """Creates a list of hole declarations from a string encoding. Example of the string encoding: "H0,((Start Bool (a b)));H1,((Start Int (x y (+ Start Start))))"."""
        return [HoleDecl.from_str(x) for x in text.split(";") if x != ""]




class SynthesisResult(solvers.SolverResult):
    """SynthesisResult contains results of the synthesis. One of the artifacts produced by this class upon creation is final synthesized source code of the program, i.e. program with filled holes.

    Attributes:
    -----------
    text : string
        Raw output of the solver.
    decision : string
        Solver decision: sat, unsat, unsupported or unknown.
    model : dict[string,string]
        Dictionary containing values of all variables present in the model.
    original_code : string
        Original python code of the program.
    holes_defs : list[HoleDecl]
        List of hole declarations.
    used_struct_vars : list[string]
        List of structural variables used for creating final program.
    """

    def __init__(self, solver_result, original_code, holes_decls, env):
        solvers.SolverResult.__init__(self, solver_result, env)
        self.original_code = original_code
        self.holes_decls = holes_decls
        self.used_struct_vars = []
        self.holes_content = self.get_holes_content()
        self.final_code = self.get_final_code()

    def get_holes_content(self):
        """Returns a map containing synthesized content of the hole in the SMT-LIB 2.0 language."""
        if self.decision != 'sat':
            return {}
        else:
            res = {}
            for hole in self.holes_decls:
                res[hole.id] = self.get_source_code_for_hole(hole)
            return res

    def get_final_code(self):
        """Produces code of the synthesized program. Assumes that every name in the program is unique, so that standard
        string replace method may be used."""
        if self.decision != 'sat':
            return self.original_code

        final_code = self.original_code
        for hole in self.holes_decls:
            final_code = final_code.replace(hole.id, self.get_source_code_for_hole(hole))

        if self.env.synth_substitute_free:  # By default true
            # TODO: Substitute only explicitly given free variables and not all vars from the model.
            # It works only because non-free vars are iterated on in forall, and in the case of synthesis from test cases
            # all non-free variables are renamed (so there will be no collision with the names in the original code).
            for name in self.model:
                final_code = final_code.replace(name, self.model[name])
        return final_code

    def get_source_code_for_hole(self, hole):
        """Converts solver output of a given hole into Python source code."""
        gt = hole.grammar_tree
        smt = self.get_merged_hole_result(gt.root)
        return self.convert_to_python(smt)

    def get_merged_hole_result(self, gra):
        """Returns SMTLIB formula based on production rules chosen by solver."""
        assert type(gra) == GrammarRuleApplier

        self.used_struct_vars.append(gra.struct_var_name)
        if gra.struct_var_name not in self.model:
            # Sometimes Z3 ignores values of certain variables in the model if they are not necessary.
            # In such a case, an id of a hole is returned.
            return gra.hole_decl.id

        val = int(self.model[gra.struct_var_name])
        node = gra[val]
        if node.is_leaf:
            return self.substitute_const_variables(node.get_text(), node)
        else:
            ac = 0
            rsymbs = node.rule_symbols
            node_body = node.body[:]
            for r in rsymbs:
                pos = node_body.index(r)
                node_body[pos] = self.get_merged_hole_result(node.args[ac])
                ac += 1
            node_text = Grammar.body_text(node_body)
            return self.substitute_const_variables(node_text, node)

    def substitute_const_variables(self, node_text, node):
        assert isinstance(node, GrammarNode)
        for v_name in node.vars:
            value = self.model[v_name]
            node_text = node_text.replace(v_name, value)
        return node_text

    def convert_to_python(self, smt):
        return smt



class PythonSMTLIB():
    py2smt = {}
    smt2py = {}

    def __init__(self):
        # Arithmetic
        PythonSMTLIB.add('+', '+')
        PythonSMTLIB.add('-', '-')
        PythonSMTLIB.add('*', '*')
        PythonSMTLIB.add('/', '/')
        PythonSMTLIB.add('mod', '%')
        # Logic
        PythonSMTLIB.add('and', 'and')
        PythonSMTLIB.add('or', 'or')
        PythonSMTLIB.add('not', 'not')

    @staticmethod
    def add(smt, py):
        PythonSMTLIB.py2smt[py] = smt
        PythonSMTLIB.smt2py[smt] = py

