kodkod_annotator

annotate_constraint_for_kodkod(Ast, VarsAcc, [CoupledVars], AnnotatedAst) :-
    AstAsBool = b(convert_bool(Ast),boolean,[]),
    find_typed_identifier_uses(Ast, [], LocalIds),
    join_using_couples(LocalIds, CoupledLocalIds),
    join_using_couples([CoupledLocalIds|VarsAcc], CoupledVars),
    EmptySet = b(empty_set,set(any),[]),
    AnnotatedAst = b(external_pred_call('KODKOD',[b(integer(1),integer,[]), CoupledLocalIds, EmptySet, AstAsBool]),pred,[]).

annotate_kodkod_constraints(_, Ast, AnnotatedAst) :-
    Ast = b(kodkod(_,_),_,_),
    !,
    AnnotatedAst = Ast.

annotate_kodkod_constraints(AnnotationType, Ast, AnnotatedAst) :-
    annotate_kodkod_constraints(AnnotationType, Ast, [], NewVarsAcc, TAnnotatedAst),
    NewVarsAcc = [CoupledVars],
    EmptySet = b(empty_set,set(any),[]),
    KodkodSolve = b(external_pred_call('KODKOD_SOLVE',[b(integer(1),integer,[]), CoupledVars, EmptySet]),pred,[]),
    AnnotatedAst = b(conjunct(TAnnotatedAst, KodkodSolve),pred,[]).

annotate_kodkod_constraints(_, Ast, Ast).

annotate_kodkod_constraints(greedy, b(Node,Type,Info), VarsAcc, NewVarsAcc, AnnotatedAst) :-
    functor(Node, _, Arity),
    Arity == 2,
    arg(1, Node, Lhs),
    arg(2, Node, Rhs),
    \+contains_higher_order_or_int_ast([Lhs]),
    \+contains_higher_order_or_int_ast([Rhs]),
    !,
    annotate_constraint_for_kodkod(b(Node,Type,Info), VarsAcc, NewVarsAcc, AnnotatedAst).

annotate_kodkod_constraints(preferred, b(Node,Type,Info), VarsAcc, NewVarsAcc, AnnotatedAst) :-
    functor(Node, Functor, _),
    preferred_by_kodkod(Functor),
    !,
    annotate_constraint_for_kodkod(b(Node,Type,Info), VarsAcc, NewVarsAcc, AnnotatedAst).

annotate_kodkod_constraints(AnnotationType, b(Node,Type,Info), VarsAcc, NewVarsAcc, AnnotatedAst) :-
    functor(Node, Functor, Arity),
    Arity == 2,
    arg(1, Node, Lhs),
    arg(2, Node, Rhs),
    !,
    annotate_kodkod_constraints(AnnotationType, Lhs, VarsAcc, VarsAcc1, ALhs),
    annotate_kodkod_constraints(AnnotationType, Rhs, VarsAcc1, NewVarsAcc, ARhs),
    NNode =.. [Functor, ALhs, ARhs],
    AnnotatedAst = b(NNode,Type,Info).

annotate_kodkod_constraints(greedy, b(forall(Ids, Lhs, Rhs),pred,Info), VarsAcc, NewVarsAcc, AnnotatedAst) :-
    \+contains_higher_order_or_int_ast([Ids]),
    \+contains_higher_order_or_int_ast([Lhs]),
    \+contains_higher_order_or_int_ast([Rhs]),
    !,
    annotate_kodkod_constraints(greedy, Lhs, VarsAcc, VarsAcc1, ALhs),
    annotate_kodkod_constraints(greedy, Rhs, VarsAcc1, NewVarsAcc, ARhs),
    AnnotatedAst = b(forall(Ids, ALhs, ARhs),pred,Info).

annotate_kodkod_constraints(greedy, b(exists(Ids, Pred),pred,Info), VarsAcc, NewVarsAcc, AnnotatedAst) :-
    \+contains_higher_order_or_int_ast([Ids]),
    \+contains_higher_order_or_int_ast([Pred]),
    !,
    annotate_kodkod_constraints(greedy, Pred, VarsAcc, NewVarsAcc, APred),
    AnnotatedAst = b(exists(Ids, APred),pred,Info).

annotate_kodkod_constraints(_, Ast, VarsAcc, VarsAcc, Ast).

contains_higher_order_or_int_ast([b(Node,Type,Info)|_]) :-
    get_texpr_type(b(Node,Type,Info), Type),
    ( Type == integer
    ; Type = set(A),
    set_or_seq(A)
    ; Type = seq(A),
    set_or_seq(A)
    ),
    !.

contains_higher_order_or_int_ast([Ast|T]) :-
    Ast =.. [_|Args],
    ( contains_higher_order_or_int_ast(Args),
    !
    ; contains_higher_order_or_int_ast(T)).

join_using_couples([Id|T], Couples) :-
    join_using_couples(T, Id, Couples).

join_using_couples([], Couples, Couples).

join_using_couples([Id|T], Acc, Couples) :-
    get_texpr_type(Id, Type1),
    get_texpr_type(Acc, Type2),
    NewAcc = b(couple(Id, Acc), couple(Type1, Type2), []),
    join_using_couples(T, NewAcc, Couples).

preferred_by_kodkod(equal).

preferred_by_kodkod(not_equal).

preferred_by_kodkod(not_member).

preferred_by_kodkod(member).

preferred_by_kodkod(subset).

preferred_by_kodkod(not_subset).

preferred_by_kodkod(subset_strict).

preferred_by_kodkod(not_subset_strict).

preferred_by_kodkod(exists).

preferred_by_kodkod(forall).

preferred_by_kodkod(pow_subset).

preferred_by_kodkod(pow_subset1).

preferred_by_kodkod(fin_subset).

preferred_by_kodkod(fin_subset1).

set_or_seq(set(_)).

set_or_seq(seq(_)).