predicate_data_generator

ast_uses_id_from_list(UsedIds, Ast) :-
    Ast = b(Node,_,_),
    functor(Node, Functor, _),
    \+ blacklist_component(Functor),
    bsyntaxtree:find_identifier_uses(Ast, [], AstUsedIds),
    !,
    intersect(UsedIds, AstUsedIds).

augment_records(5).

blacklist_component(int_set).

blacklist_component(integer_set).

blacklist_component(nat_set).

blacklist_component(nat1_set).

blacklist_component(natural_set).

blacklist_component(natural1_set).

blacklist_component(forall).

blacklist_component(exists).

blacklist_component(comprehension_set).

blacklist_component(total_function).

blacklist_component(total_surjection).

blacklist_component(total_injection).

blacklist_component(total_relation).

blacklist_component(total_surjection_relation).

blacklist_component(surjection_relation).

blacklist_component(partial_function).

blacklist_component(partial_injection).

blacklist_component(partial_surjection).

blacklist_component(partial_bijection).

blacklist_component(lambda).

blacklist_component(quantified_union).

blacklist_component(quantified_intersection).

exclude_solution([], ExclusionPred, ExclusionPred) :-
    !.

exclude_solution(State, ExclusionPred, NewExclusionPred) :-
    create_equality_nodes_from_example(State, EqualityNodes),
    conjunct_predicates(EqualityNodes, EqualityConj),
    truth_or_exclude(EqualityConj,ExclusionPred, NewExclusionPred).

filter_machine_var_states([], _, Acc, Acc).

filter_machine_var_states([VarState|T], MachineVars, Acc, States) :-
    get_texpr_info(VarState, Info),
    member(synthesis(machinevar, _VarName), Info),
    % operation parameters are also marked as machinevar but not in MachineVars
    %member(b(identifier(VarName), _, _), MachineVars),
    !,
    filter_machine_var_states(T, MachineVars, [VarState|Acc], States).

filter_machine_var_states([_|T], MachineVars, Acc, States) :-
    filter_machine_var_states(T, MachineVars, Acc, States).

filter_predicate(UsedIds, Ast, NewAst) :-
    Ast = b(Node,_,_),
    Node =.. [Functor, Lhs, Rhs],
    ( Functor == conjunct
    ; Functor == disjunct
    ; Functor == equivalence
    ; Functor == implication),
    filter_predicate_binary(UsedIds, Ast, Functor, Lhs, Rhs, Clean),
    !,
    NewAst = Clean.

filter_predicate(UsedIds, Ast, NewAst) :-
    Ast = b(Node,_,_),
    Node =.. [_, Lhs, Rhs],
    get_used_components(Lhs, [], LhsUsedComponents),
    get_used_components(Rhs, [], RhsUsedComponents),
    findall(BlacklistCmpt, (member(BlacklistCmpt, LhsUsedComponents), blacklist_component(BlacklistCmpt)), LhsBlacklistCmpts),
    findall(BlacklistCmpt, (member(BlacklistCmpt, RhsUsedComponents), blacklist_component(BlacklistCmpt)), RhsBlacklistCmpts),
    !,
    ( LhsBlacklistCmpts == [],
    RhsBlacklistCmpts == []
    -> ast_uses_id_from_list(UsedIds, Ast),
    NewAst = Ast
    ; fail
).

filter_predicate(UsedIds, Negation, NegNewAst) :-
    Negation = b(negation(Ast),pred,_),
    !,
    filter_predicate(UsedIds, Ast, NewAst),
    NegNewAst = b(negation(NewAst),pred,[]).

filter_predicate(UsedIds, Ast, NewAst) :-
    ast_uses_id_from_list(UsedIds, Ast),
    get_used_components(Ast, [], UsedComponents),
    findall(BlacklistCmpt, (member(BlacklistCmpt, UsedComponents), blacklist_component(BlacklistCmpt)), BlacklistCmpts),
    BlacklistCmpts == [],
    NewAst = Ast.

filter_predicate_binary(UsedIds, b(_,Type,Info), Functor, Lhs, Rhs, NewNode) :-
    % if lhs would be empty (or truth) keep rhs only and vice versa
    ( filter_predicate(UsedIds, Lhs, NLhs)
    -> ( filter_predicate(UsedIds, Rhs, NRhs)
    -> NNode =.. [Functor, NLhs, NRhs],
    NewNode = b(NNode,Type,Info)
    ; NewNode = NLhs
    )
    ; ( filter_predicate(UsedIds, Rhs, NRhs)
    -> NewNode = NRhs
    % fail if both sides do not use any ids from UsedIds (predicate would be truth)
    ; fail
    )
    ),!.

generate_synthesis_data_from_predicate_ast(MachinePath, AugmentRecords, SolverTimeoutMs, PredicateAst, AugmentedSetOfData) :-
    preferences:temporary_set_preference(optimize_ast, false),
    preferences:temporary_set_preference(normalize_ast, true),
    % should we also set normalize_ast_sort_commutative to true?
    load_b_machine_if_unloaded(MachinePath),
    \+ current_machine_uses_records,
    retractall(tools:id_counter(_)),
    set_desired_preferences(SolverTimeoutMs, OldKodkodPref, OldTimeOutPref, OldRandPref),
    b_get_machine_variables(MachineVars),
    maplist(get_id_name, MachineVars, MachineVarNames),
    % following includes constraints not using machine variables
    %b_get_all_used_identifiers(AllUsedIds),
    b_ast_cleanup:clean_up_pred(PredicateAst, _, CleanPred),
    b_simplifier:simplify_b_predicate(CleanPred, SimplifiedPred),
    filter_predicate(MachineVarNames, SimplifiedPred, FilteredPredicateAst),
    get_library_components_from_pred_or_expr(FilteredPredicateAst, UsedComponents),
    get_normalized_id_name_mapping_stateful(NormalizedSets, NormalizedIds, NOperationNames, NRecordFieldNames),
    Env = [[], NormalizedSets, NormalizedIds, NOperationNames, NRecordFieldNames],
    %b_get_typed_invariant_from_machine(Invariant),
    find_typed_identifier_uses(FilteredPredicateAst, [], UsedIds),
    translate:generate_typing_predicates(UsedIds, TypingPredicates),
    conjunct_predicates(TypingPredicates, TypingPredicate),
    %PredWithTypes = b(conjunct(TypingPredicate, FilteredPredicateAst),pred,[]),
    get_augmented_states_for_predicate(Env, AugmentRecords, MachineVars, UsedComponents, FilteredPredicateAst, TypingPredicate, AugmentedSetOfData, _),
    reset_old_preferences(OldKodkodPref, OldTimeOutPref, OldRandPref),
    preferences:reset_temporary_preference(optimize_ast),
    preferences:reset_temporary_preference(normalize_ast).

generate_synthesis_data_from_predicate_raw(MachinePath, RawPredicate, AugmentedSetOfData) :-
    augment_records(AR),
    solver_timeout_ms(SolverTimeoutMs),
    generate_synthesis_data_from_predicate_raw(MachinePath, AR, SolverTimeoutMs, RawPredicate, AugmentedSetOfData).

generate_synthesis_data_from_predicate_raw(MachinePath, AugmentRecords, SolverTimeoutMs, RawPredicate, AugmentedSetOfData) :-
    atom_codes(RawPredicate, RawPredicateCodes),
    parse_formula(RawPredicateCodes, UntypedPredAst),
    generate_synthesis_data_from_predicate_untyped(MachinePath, AugmentRecords, SolverTimeoutMs, UntypedPredAst, AugmentedSetOfData).

generate_synthesis_data_from_predicate_untyped(MachinePath, AugmentRecords, SolverTimeoutMs, UntypedPredAst, AugmentedSetOfData) :-
    load_b_machine_if_unloaded(MachinePath),
    type_check_in_machine_context([UntypedPredAst], TypedPredAsts),
    TypedPredAsts = [TypedPredAst],
    generate_synthesis_data_from_predicate_ast(MachinePath, AugmentRecords, SolverTimeoutMs, TypedPredAst, AugmentedSetOfData).

get_amount_of_states_for_predicate(Env, 0, _, _, _, Acc, Acc, Env) :-
    !.

get_amount_of_states_for_predicate(Env, AmountOfExamples, ExclusionPred, MachineVars, PredicateAst, Acc, ListOfExamples, NewEnv) :-
    AmountOfExamples1 is AmountOfExamples-1,
    solve_predicate(b(conjunct(ExclusionPred,PredicateAst),pred,[]), _, 1, [force_evaluation], Solution),
    Solution = solution(Bindings),
    get_input_nodes_from_bindings(Bindings, TempState),
    filter_machine_var_states(TempState, MachineVars, [], State),
    State \== [],
    exclude_solution(State, ExclusionPred, NewExclusionPred),
    map_translate_var_state(Env, State, PrettyState, Env1),
    !,
    get_amount_of_states_for_predicate(Env1, AmountOfExamples1, NewExclusionPred, MachineVars, PredicateAst, [PrettyState|Acc], ListOfExamples, NewEnv).

get_amount_of_states_for_predicate(Env, _, _, _, _, Acc, Acc, Env).

get_augmented_states_for_predicate(Env, AR, MachineVars, UsedComponents, PredicateAst, TypingPredicate, AugmentedSetOfData, NewEnv) :-
    random_list_of_numbers(AR, [], RandomNrList),
    get_augmented_states_for_predicate_rand(Env, RandomNrList, MachineVars, UsedComponents, PredicateAst, TypingPredicate, AugmentedSetOfData, NewEnv).

get_augmented_states_for_predicate_rand(Env, [], _, _, _, _, [], Env).

get_augmented_states_for_predicate_rand(Env, [AmountOfExamples|T], MachineVars, UsedComponents, PredicateAst, TypingPredicate, [(PositiveStates,NegativeStates,UsedComponents)|NT], NewEnv) :-
    AmountOfExamples = (PAmountOfExamples,NAmountOfExamples),
    Pos = b(conjunct(TypingPredicate,PredicateAst),pred,[]),
    Neg = b(conjunct(TypingPredicate,b(negation(PredicateAst),pred,[])),pred,[]),
    get_amount_of_states_for_predicate(Env, PAmountOfExamples, b(truth,pred,[]), MachineVars, Pos, [], PositiveStates, Env1),
    get_amount_of_states_for_predicate(Env1, NAmountOfExamples, b(truth,pred,[]), MachineVars, Neg, [], NegativeStates, Env2),
    length(PositiveStates, LPos),
    length(NegativeStates, LNeg),
    Amount is LPos+LNeg,
    min_amount_of_examples(MinAmountOfExamples),
    Amount >= MinAmountOfExamples,
    !,
    get_augmented_states_for_predicate_rand(Env2, T, MachineVars, UsedComponents, PredicateAst, TypingPredicate, NT, NewEnv).

get_augmented_states_for_predicate_rand(Env, [_|T], MachineVars, UsedComponents, PredicateAst, TypingPredicate, NT, NewEnv) :-
    get_augmented_states_for_predicate_rand(Env, T, MachineVars, UsedComponents, PredicateAst, TypingPredicate, NT, NewEnv).

get_id_name(b(identifier(Name),_,_),Name).

get_normalized_id_name_mapping_stateful(NormalizedSets, NormalizedIds, NOperationNames, NRecordFieldNames) :-
    normalized_id_name_mapping(NormalizedSets, NormalizedIds, NOperationNames, NRecordFieldNames).

get_normalized_id_name_mapping_stateful(NormalizedSets, NormalizedIds, NOperationNames, NRecordFieldNames) :-
    \+ normalized_id_name_mapping(_, _, _, _),
    get_normalized_id_name_mapping(NormalizedSets, NormalizedIds, NOperationNames, NRecordFieldNames),
    asserta(normalized_id_name_mapping(NormalizedSets, NormalizedIds, NOperationNames, NRecordFieldNames)).

get_random_amount_of_examples((PAmountOfExamples,NAmountOfExamples)) :-
    min_amount_of_examples(MinAmountOfExamples),
    max_amount_of_examples(positive, PTempMaxAmountOfExamples),
    max_amount_of_examples(negative, NTempMaxAmountOfExamples),
    PMaxAmountOfExamples is PTempMaxAmountOfExamples+1,
    NMaxAmountOfExamples is NTempMaxAmountOfExamples+1,
    random(MinAmountOfExamples, PMaxAmountOfExamples, PAmountOfExamples),
    random(MinAmountOfExamples, NMaxAmountOfExamples, NAmountOfExamples).

get_used_components(b(Node,_,_), Acc, [Functor|Acc]) :-
    functor(Node, Functor, 0),!.

get_used_components(b(Node,_,_), Acc, UsedComponents) :-
    Node =.. [Component|Args],
    map_get_used_components(Args, [Component|Acc], UsedComponents),!.

get_used_components(AtomNr, Acc, Acc) :-
    functor(AtomNr, _, 0).

load_b_machine_if_unloaded(MachinePath) :-
    bmachine_is_precompiled,
    b_get_main_filename(LoadedMachinePath),
    atom_concat(_, MachinePath, LoadedMachinePath),
    !.

load_b_machine_if_unloaded(MachinePath) :-
    retractall(normalized_id_name_mapping(_, _, _, _)),
    b_load_machine_from_file(MachinePath),
    b_machine_precompile.

map_get_used_components([], Acc, Acc).

map_get_used_components([Arg|T], Acc, UsedComponents) :-
    get_used_components(Arg, Acc, NewAcc),
    map_get_used_components(T, NewAcc, UsedComponents).

map_translate_var_state(Env, [], [], Env).

map_translate_var_state(Env, [Ast|T], [(VarName,Type,PrettyAst)|NT], NewEnv) :-
    normalize_ids_in_b_ast(Env, Ast, NAst, Env1),
    NAst = b(_, Type, Info),
    member(synthesis(machinevar, VarName), Info),
    translate_bexpression(NAst, PrettyAst),
    map_translate_var_state(Env1, T, NT, NewEnv).

max_amount_of_examples(positive, 8).

max_amount_of_examples(negative, 8).

min_amount_of_examples(3).

random_list_of_numbers(0, Acc, Acc) :-
    !.

random_list_of_numbers(C, Acc, L) :-
    C1 is C-1,
    get_random_amount_of_examples(R),
    \+ member(R, Acc),
    !,
    random_list_of_numbers(C1, [R|Acc], L).

random_list_of_numbers(C, Acc, L) :-
    random_list_of_numbers(C, Acc, L).

reset_old_preferences(OldKodkodPref, OldTimeOutPref, OldRandPref) :-
    set_preference(try_kodkod_on_load, OldKodkodPref),
    set_preference(time_out, OldTimeOutPref),
    set_preference(randomise_enumeration_order, OldRandPref).

set_desired_preferences(Timeout, OldKodkodPref, OldTimeOutPref, OldRandPref) :-
    get_preference(try_kodkod_on_load, OldKodkodPref),
    set_preference(try_kodkod_on_load, false),
    get_preference(randomise_enumeration_order, OldRandPref),
    set_preference(randomise_enumeration_order, true),
    get_preference(time_out, OldTimeOutPref),
    set_preference(time_out, Timeout).

solver_timeout_ms(10000).

truth_or_exclude(b(truth,pred,[]), ExclusionPred, ExclusionPred) :-
    !.

truth_or_exclude(EqualityConj, ExclusionPred, b(conjunct(b(negation(EqualityConj),pred,[]),ExclusionPred),pred,[])).