memoization

apply_to_memoize(MemoID,P,T,B,X,Y,Span,WF) :-
    ground_value_check(X,XV),
    block_apply_closure_memo_groundx(XV,MemoID,X,Y,P,T,B,raise_wd,Span,WF).

assert_stored_memo_result_id(1,ArgHash,Argument,StoredResult) :- !,
    assert(stored_memo_result1(ArgHash,Argument,StoredResult)).

assert_stored_memo_result_id(2,ArgHash,Argument,StoredResult) :- !,
    assert(stored_memo_result2(ArgHash,Argument,StoredResult)).

assert_stored_memo_result_id(3,ArgHash,Argument,StoredResult) :- !,
    assert(stored_memo_result3(ArgHash,Argument,StoredResult)).

assert_stored_memo_result_id(4,ArgHash,Argument,StoredResult) :- !,
    assert(stored_memo_result4(ArgHash,Argument,StoredResult)).

assert_stored_memo_result_id(5,ArgHash,Argument,StoredResult) :- !,
    assert(stored_memo_result5(ArgHash,Argument,StoredResult)).

assert_stored_memo_result_id(6,ArgHash,Argument,StoredResult) :- !,
    assert(stored_memo_result6(ArgHash,Argument,StoredResult)).

assert_stored_memo_result_id(7,ArgHash,Argument,StoredResult) :- !,
    assert(stored_memo_result7(ArgHash,Argument,StoredResult)).

assert_stored_memo_result_id(8,ArgHash,Argument,StoredResult) :- !,
    assert(stored_memo_result8(ArgHash,Argument,StoredResult)).

assert_stored_memo_result_id(MemoID,ArgHash,Argument,StoredResult) :-
    assert(stored_memo_result(ArgHash,MemoID,Argument,StoredResult)).

bb_reset_reuse_counter(_).

block_apply_closure_memo_groundx(_,MemoID, X,Y, P,T,B, RaiseWD, Span,WF) :-
    % MemoID is ground, meaning body B is ground as well
    % TODO: obtain stored value directly
    compute_funcall_memo_hash(X,MemoID,ArgHash),
    if(get_stored_memo_result(MemoID,X,ArgHash,StoredResult),
    equal_object_wf(Y,StoredResult,WF),
    (start_profile_function_application(MemoID,Timer),
    if(check_element_of_function_closure_nowf(MemoID,X,ComputedResult,P,T,B),
    % ground WF to compute ground value for ComputedResult
    (stop_profile_function_application(MemoID,Timer),
    (ground_value(ComputedResult)
    -> store_memo_result(MemoID,X,ArgHash,ComputedResult,Span,WF)
    ; add_error(memoization,'Could not compute ground value for function application',MemoID,Span),
    print('Argument: '),translate:print_bvalue(X),nl,
    print('Result: '), print(ComputedResult),nl,
    print('Function: '),registered_value(MemoID,_Name,FunValue), translate:print_bvalue(FunValue),nl,nl
    ),
    % To do: check if we compute two or more solutions for f(X) and add WD error
    equal_object_wf(Y,ComputedResult,WF)
    ),
    (RaiseWD=raise_wd,
    add_wd_error_span('function applied outside of domain (#8-Memo): ', '@fun'(X,closure(P,T,B)),Span,WF)
)
)
)
).

check_element_of_function_closure_nowf(MemoID,X,ComputedResult,_,_,_) :-
    registered_value(MemoID,_,closure(P,T,B)),!,
    init_wait_flags(WF,[check_element_of_function_closure_nowf(MemoID)]),
    check_element_of_function_closure(X,ComputedResult,P,T,B,WF),
    ground_wait_flags(WF).

check_element_of_function_closure_nowf(_,X,ComputedResult,P,T,B) :-
    init_wait_flags(WF,[check_element_of_function_closure_nowf]),
    check_element_of_function_closure(X,ComputedResult,P,T,B,WF),
    ground_wait_flags(WF).

compile_value(_,_).

compute_funcall_memo_hash(int(X),MemoID,R) :- MemoID <9, !,
    R=X. % avoid computing hash for simple integer values; only drawback: indexing not good for too large integers

compute_funcall_memo_hash(X,MemoID,ArgHash) :-
    compute_memo_hash((MemoID,X),ArgHash).

compute_memo_hash(LookupTerm,Hash) :-
    inc_counter(memoization_hash_counter,_),
    term_hash(LookupTerm,
    [range(smallint),algorithm(sdbm), depth(infinite),if_var(ignore)],Hash).

construct_memoization_closure(MemoID,BValue,ElType,Result) :-
    get_infos(BValue,BInfos),
    XFC = b(external_function_call('MEMOIZE_STORED_FUNCTION',[b(value(int(MemoID)),integer,[])]),set(ElType),[]),
    TID = b(identifier('_zzzz_unary'),ElType,[]),
    Result = closure(['_zzzz_unary'],[ElType],
    b(member(TID,XFC),pred,
    [prob_annotation('MEMOIZE'(MemoID))|BInfos])).

count_reuse(_MemoID) :-
    inc_counter(memoization_function_reuse_counter,_).

element_of_memoization_closure(MemoID,Element,WF,P,T,B) :-
    nonvar(Element), Element = (X,Y),
    ground_value(X),
    % we can do a function application:
    !, %print(fun_apply_reuse(MemoID,Element)),nl,
    Span=B,
    block_apply_closure_memo_groundx(ground,MemoID, X,Y, P,T,B, do_not_raise_wd, Span,WF).

element_of_memoization_closure(MemoID,Element,WF,_,_,B) :-
    % TO DO: delay before computing complete memoization
    get_complete_memoization_expansion(MemoID,Set,_Done,B,element_of_memoization_closure,WF),
    !, %print(el_of(MemoID,Element)),nl,
    check_element_of_wf(Element,Set,WF).

expanded_value(avl_set(_)).

expanded_value(global_set(_)).

expanded_value([]).

expanded_value([_|_]).

function_application_time_available :- fail.

get_complete_memoization_expansion(MemoID,Result,Done,_,_PP,WF) :-
    try_get_complete_memoization_expansion(MemoID,Cached,WF),
    !,
    Result=Cached, Done=true.

get_complete_memoization_expansion(MemoID,Result,Done,Span,PP,WF) :-
    registered_value(MemoID,Name,Val), nonvar(Val),
    !,
    get_complete_memoization_expansion_aux(Val,MemoID,Name,Result,Done,Span,PP,WF).

get_complete_memoization_expansion(MemoID,_,_,Span,PP,_) :-
    add_error(get_complete_memoization_expansion,'Memoized function ID not registered:',MemoID:PP,Span),
    fail.

get_complete_memoization_expansion_aux(Val,_MemoID,_,Result,Done,_Span,_PP,_WF) :- expanded_value(Val), !,
    Result=Val,Done=true.

get_complete_memoization_expansion_aux(closure(P,T,B),MemoID,Name,Result,Done,Span,PP,WF) :- !,
    debug_println(19,expanding_memoized_function(MemoID,PP)),
    (custom_explicit_sets:is_infinite_or_very_large_closure(P,T,B,200000)
    -> add_warning(memoization(PP),'Expanding infinite or very large memoized function: ',Name,B)
    ; custom_explicit_sets:is_symbolic_closure(P,T,B)
    -> add_message(memoization(PP),'Fully expanding symbolic memoized function: ',Name,B)
    ; true),
    custom_explicit_sets:expand_normal_closure_direct(P,T,B,Expansion,DoneExp,WF),
    store_memoization_closure_expansion(MemoID,DoneExp,Expansion,Done,Result,Span,WF).

get_complete_memoization_expansion_aux(Val,MemoID,Name,_,_,Span,PP,_) :-
    add_error(memoization(PP),'Illegal value stored for memoized function:',MemoID:Name,Span),
    print(Val),nl,
    fail.

get_infos(Var,Res) :- var(Var),!,Res=[].

get_infos(closure(_,_,B),Res) :- extract_info(B,Infos),!,
    get_texpr_info(B,AllInfos),extract_pos_infos(AllInfos,Pos),
    append(Pos,Infos,Res).

get_infos(_,[]). % :- print(no_infos(A)),nl.

get_memoization_closure_value(P,T,B,Value) :-
    is_memoization_closure(P,T,B,MemoID),
    number(MemoID),
    registered_value(MemoID,_,Value).

get_registered_function(Id,Val) :- registered_value(Id,_,Val).

get_stored_memo_expansion(Hash,Term,Result) :-
    stored_memo_expansion(Hash,Term,Result),
    inc_counter(memoization_reuse_counter,_).

get_stored_memo_result(MemoID,Argument,ArgHash,StoredResult) :-
    stored_memo_result_id(MemoID,ArgHash,Argument,StoredResult),
    %print(get_stored(MemoID,ArgHash,StoredResult)),nl,
    count_reuse(MemoID).

inc(Hash) :- (retract(hash_count(Hash,Nr)) -> N1 is Nr+1 ; N1 = 1), assert(hash_count(Hash,N1)).

init_memo_table :-
    counter_init,
    new_counter(memoization_reuse_counter), new_counter(memoization_hash_counter),
    new_counter(memoization_function_counter), new_counter(memoization_function_reuse_counter),
    reset_memo_table.

is_memoization_closure([ID],_T,b(Body,_,INFO),MemoID) :-
    member(prob_annotation('MEMOIZE'(MemoID)),INFO),
    % check Body has not been rewritten:
    Body = member(TID,b(RHS,_,_)),
    get_texpr_id(TID,ID),
    is_memoize_stored_function_reference(RHS,MemoID).

is_memoize_stored_function_reference(external_function_call('MEMOIZE_STORED_FUNCTION',[Arg]),MemoID) :-
    Arg=b(value(int(MemoID)),_,_).

my_avl_domain(avl_set(A),List) :- avl_domain(A,List).

my_avl_domain([],[]).

my_bb_get_reuse_counter(_,'?').

not_element_of_memoization_closure(MemoID,Element,WF,P,T,B) :-
    nonvar(Element), Element = (X,Y),
    ground_value(X),
    !,
    if(block_apply_closure_memo_groundx(ground,MemoID, X,YRes, P,T,B, do_not_raise_wd, B,WF),
    not_equal_object_wf(Y,YRes,wf),
    true).

not_element_of_memoization_closure(MemoID,Element,WF,_,_,B) :-
    % TO DO: delay before computing complete memoization
    get_complete_memoization_expansion(MemoID,Set,_Done,B,not_element_of_memoization_closure,WF),
    !,
    not_element_of_wf(Element,Set,WF).

print_memo_entry(closure(Parameters,ParameterTypes,ClosureBody)) :- !,
    translate:print_bvalue(closure(Parameters,ParameterTypes,ClosureBody)).

print_memo_entry(closure1_for_explicit_set(A)) :- !,
    print('closure1('),translate:print_bvalue(avl_set(A)), print(')').

print_memo_entry(A) :- translate:print_bvalue(A).

print_memo_profile :-
    ((stored_memo_expansion(_,_,_) ; registered_value(_,_,_))
    -> print('MEMO Table:'),nl,
    print_memo_table
    ; true).

print_memo_table :-
    stored_memo_expansion(Hash,LookupTerm,PackedValue),
    (hash_count(Hash,Count) -> print(Count), print(' : ') ; true),
    (hash_time(Hash,Time) -> print(Time),print(' --> ') ; true),
    print_memo_entry(LookupTerm),
    nl, print(' : '),translate:print_bvalue(PackedValue),nl,fail.

print_memo_table :-
    debug_mode(on),
    registered_value(MemoID,Name,FunValue),
    (memoization_expansion_completed(MemoID)
    -> Compl=' (completely evaluated)' ; Compl=''),
    format('MemoID ~w stored FUNCTION ~w~w ',[MemoID,Name,Compl]),
    nl, print(' : '),translate:print_bvalue(FunValue),nl,fail.

print_memo_table :-
    debug_mode(on),
    registered_value(MemoID,Name,_),
    stored_memo_result_id(MemoID,_Hash,LookupTerm,PackedValue),
    format('MemoID ~w (~w) result for argument ',[MemoID,Name]),
    print_memo_entry(LookupTerm),
    nl, print(' : '),translate:print_bvalue(PackedValue),nl,fail.

print_memo_table :-
    stored_memo_result(Hash,MemoID1,T1,_),
    stored_memo_result(Hash,MemoID2,T2,_),
    dif((MemoID1,T1),(MemoID2,T2)),
    format('Hash COLLISION ~w, ~w:~w vs ~w:~w~n',[Hash,MemoID1,T1,MemoID2,T2]),
    fail.

print_memo_table :- format('Summary of reuse per memoization ID:~n',[]),
    registered_value(MemoID,Name,_),
    my_bb_get_reuse_counter(MemoID,Cnt),
    findall(1,stored_memo_result_id(MemoID,_Hash,_,_),Stored), length(Stored,NrStored),
    (\+ function_application_time_available -> WTime='?'
    ; function_application_time(MemoID,WTime) -> true ; WTime=0),
    format('MemoID ~w (~w) : Values Stored: ~w (~w ms to compute), Reused ~w~n',[MemoID,Name,NrStored,WTime,Cnt]),
    fail.

print_memo_table :- get_counter(memoization_hash_counter,Hash),
    get_counter(memoization_reuse_counter,RHit),
    format('Hashes computed: ~w, expansions reused: ~w~n',[Hash,RHit]),
    get_counter(memoization_function_counter,F),
    get_counter(memoization_function_reuse_counter,FR),
    format('Memoization functions registered: ~w, results reused: ~w~n',[F,FR]).

register_memoization_function(Name,BValue,ElType,RecursionInfo,MemoID,MemoResultClosure) :-
    value_variables(BValue,Vars),
    %print(vars(Name,Vars,LHSValue)),nl, print(bvalue(BValue)),nl,
    (RecursionInfo='$recursion_value'(RecValue)
    -> remove_variables(Vars,[RecValue],WVars), % do not wait on RecValue; it will be set by registering MEMOID
    RecValue = MemoResultClosure,
    InnerMemoID=MemoID % already instantiate MemoID before asserting closure; cyclic use
    ; WVars=Vars
    ),
    ground_value_check(WVars,Ground),
    %print(register(Name,MemoID,Ground,BValue)),nl,nl,
    construct_memoization_closure(MemoID,BValue,ElType,MemoResultClosure),
    register_memoization_function_aux(Ground,Name,BValue,InnerMemoID,MemoID).

register_memoization_function_aux(_,Name,BValue,InnerMemoID,MemoID) :-
    inc_counter(memoization_function_counter,InnerMemoID),
    format('Registering function ~w for memoization: ~w~n',[InnerMemoID,Name]),
    (ground(BValue) -> true ; add_internal_error('Non ground memo closure: ',InnerMemoID:BValue)),
    compile_value(BValue,InnerMemoID),
    assert(registered_value(InnerMemoID,Name,BValue)),
    %if(MemoID=InnerMemoID, format('Finish ~w~n',[MemoID]),format('FAILURE ~w~n',[InnerMemoID])).
    MemoID=InnerMemoID. % only unify after assert; in case co-routines are triggered which call get_registered_function

reset_cogen.

reset_memo_table :-
    retract(registered_value(MemoID,_,_)),
    bb_reset_reuse_counter(MemoID),
    fail.

reset_memo_table :-
    reset_counter(memoization_reuse_counter),
    reset_counter(memoization_hash_counter),
    reset_counter(memoization_function_counter),
    reset_counter(memoization_function_reuse_counter),
    (stored_memo_expansion(_,_,_) -> debug_println(19,'Resetting MEMO table') ; true),
    retractall(hash_count(_,_)),
    retractall(hash_time(_,_)),
    retractall(stored_memo_expansion(_,_,_)),
    retractall(registered_value(_,_,_)),
    retractall(function_application_time(_,_)),
    retractall(memoization_expansion_completed(_)),
    retractall_stored_memo_result,
    reset_cogen.

reset_memo_with_statistics :-
    retract_with_statistics(memoization,
    [stored_memo_expansion(_,_,_),registered_value(_,_,_),memoization_expansion_completed(_),
    stored_memo_result1(_,_,_), stored_memo_result2(_,_,_),
    stored_memo_result3(_,_,_), stored_memo_result4(_,_,_),
    stored_memo_result5(_,_,_), stored_memo_result6(_,_,_),
    stored_memo_result7(_,_,_), stored_memo_result8(_,_,_),
    stored_memo_result(_,_,_,_)]),
    reset_memo_table.

retractall_stored_memo_result :-
    retractall(stored_memo_result1(_,_,_)),
    retractall(stored_memo_result2(_,_,_)),
    retractall(stored_memo_result3(_,_,_)),
    retractall(stored_memo_result4(_,_,_)),
    retractall(stored_memo_result5(_,_,_)),
    retractall(stored_memo_result6(_,_,_)),
    retractall(stored_memo_result7(_,_,_)),
    retractall(stored_memo_result8(_,_,_)),
    retractall(stored_memo_result(_,_,_,_)).

start_profile_function_application(_,0).

stop_profile_function_application(_,_).

store_function_result(MemoID,Span,WF,(X,ComputedResult)) :-
    compute_funcall_memo_hash(X,MemoID,ArgHash),
    store_memo_result(MemoID,X,ArgHash,ComputedResult,Span,WF). % TO DO: we could use special version that also allows relations

store_memo_computation_time(Hash,Time) :- % for closure expansions, not for funct. applications
    retractall(hash_time(Hash,_)), assert(hash_time(Hash,Time)).

store_memo_expansion(Hash,LookupTerm,StoredResult) :-
    assert(stored_memo_expansion(Hash,LookupTerm,StoredResult)).

store_memo_result(MemoID,Argument,ArgHash,Result,Span,WF) :-
    normalise_value_for_var(MemoID,Result,NormResult),
    store_norm_memo_result(MemoID,Argument,ArgHash,NormResult,Span,WF).

store_memoization_closure_expansion(MemoID,_,Res,OuterDone,OuterRes,Span,WF) :-
    my_avl_domain(Res,List),
    !,
    debug_println(19,storing_expansion_of_memoized_function(MemoID)),
    (maplist(store_function_result(MemoID,Span,WF),List)
    -> assert(memoization_expansion_completed(MemoID))
    ; add_internal_error(memoization,'Error storing expansion for memoized function:',MemoID,Span)),
    equal_object_wf(Res,OuterRes,WF), OuterDone=true.

store_memoization_closure_expansion(MemoID,_,R,OuterDone,OuterRes,Span,WF) :-
    add_warning(memoization,'Cannot store expansion for memoized function:',MemoID:R,Span),
    equal_object_wf(R,OuterRes,WF), OuterDone=true.

store_norm_memo_result(MemoID,Argument,ArgHash,Result,Span,WF) :-
    stored_memo_result_id(MemoID,ArgHash,Argument,StoredResult),
    !,
    (Result=StoredResult -> true
    ; equal_object(Result,StoredResult) -> format('Two representations for memoized function (~w) result.~n',[MemoID])
    ; ajoin(['Memoized function ',MemoID,' has multiple results for: '],Msg),
    Relation = [(Argument,StoredResult),(Argument,Result)],
    add_wd_error_span(Msg,'@fun'(Argument,Relation),Span,WF)
%,print('stored: '),translate:print_bvalue(StoredResult),nl,print('new: '),translate:print_bvalue(Result),nl
).

store_norm_memo_result(MemoID,Argument,ArgHash,Result,_,_) :-
    %format('Storing memo result for ~w, argument =',[MemoID]),
    %translate:print_bvalue(Argument),nl, translate:print_bvalue(Result),nl,
    assert_stored_memo_result_id(MemoID,ArgHash,Argument,Result).

stored_memo_result_id(1,ArgHash,Argument,StoredResult) :- !,
    stored_memo_result1(ArgHash,Argument,StoredResult).

stored_memo_result_id(2,ArgHash,Argument,StoredResult) :- !,
    stored_memo_result2(ArgHash,Argument,StoredResult).

stored_memo_result_id(3,ArgHash,Argument,StoredResult) :- !,
    stored_memo_result3(ArgHash,Argument,StoredResult).

stored_memo_result_id(4,ArgHash,Argument,StoredResult) :- !,
    stored_memo_result4(ArgHash,Argument,StoredResult).

stored_memo_result_id(5,ArgHash,Argument,StoredResult) :- !,
    stored_memo_result5(ArgHash,Argument,StoredResult).

stored_memo_result_id(6,ArgHash,Argument,StoredResult) :- !,
    stored_memo_result6(ArgHash,Argument,StoredResult).

stored_memo_result_id(7,ArgHash,Argument,StoredResult) :- !,
    stored_memo_result7(ArgHash,Argument,StoredResult).

stored_memo_result_id(8,ArgHash,Argument,StoredResult) :- !,
    stored_memo_result8(ArgHash,Argument,StoredResult).

stored_memo_result_id(MemoID,ArgHash,Argument,StoredResult) :-
    stored_memo_result(ArgHash,MemoID,Argument,StoredResult).

try_get_complete_memoization_expansion(MemoID,Result,WF) :-
    memoization_expansion_completed(MemoID),
    debug_println(19,recovering_full_value_for_memoized_function(MemoID)),
    findall((X,Y),stored_memo_result_id(MemoID,_,X,Y),List), % probably not sorted !
    equal_object_optimized_wf(List,Result,expand_normal_closure_memo,WF).