
?- op(700, xfx, ->).

% Phase I
% -------
%   Decorate tree with alignment map fields, instantiate domains at each
% node, and impose ``horizontal constraints'', and ``vertical constraints
% at `arr' nodes.

alignDomain(fun(F, I, D), fun(F, I, E, align(N, T), align(P, U)), P, U) :-
    daughtersDomain(D, E, L, T),
    arrExpand(E, L),
    domain(M, I),
    union(L, M, N),
    domain(N, J),
    domain(P, J).
alignDomain(arr(I, D), arr(I, E, align(T, N), align(T, _)), T, P) :-
    daughtersDomain(D, E, T, L),
    domain(M, I),
    union(L, M, N),
    subtract(L, M, P).
alignDomain(var(V, D), var(V, E, align(T, M), align(U, N)), U, N) :-
    daughtersDomain(D, E, T, M),
    domain(M, I),
    domain(N, I).
alignDomain(idx(X), idx(X, align(T, M)), T, M) :-
    domain(M, [X]).
    
daughtersDomain([], [], _, []).
daughtersDomain([X | D], [Y | E], T, N) :-
    alignDomain(X, Y, T, L),
    daughtersDomain(D, E, T, M),
    union(L, M, N).

arrExpand([], _).
arrExpand([arr(_, _, align(_, L), align(_, N)) | D], M) :-
    !,
    union(L, M, N),
    arrExpand(D, M).
arrExpand([_ | D], M) :- arrExpand(D, M).
  
domain([], []).
domain([X -> _ | M], [X | I]) :- domain(M, I).

% Phase II
% --------
%   Instantiate local-reference, synthesised or inherited mappings at
% variable and function reference nodes.  Attempt to propagate node
% mappings up to parent.

alignMapping1(Q, fun(_, _, D, align(T, M), align(U, N))) :-
    daughtersMapping1(Q, D),
    synthOrInheritMap(T, M, U, N).
alignMapping1(Q, arr(_, D, _, _)) :- daughtersMapping1(Q, D).
alignMapping1(Q, var(V, D, align(T, M), align(U, N))) :-
    localMap(Q, V, D, T, M), !,
    daughtersMapping1(Q, D),
    parentMap(T, M, U, N).
alignMapping1(Q, var(_, D, align(T, M), align(U, N))) :-
    daughtersMapping1(Q, D),
    synthOrInheritMap(T, M, U, N).
alignMapping1(_, idx(_, _)).

daughtersMapping1(_, []).
daughtersMapping1(Q, [X | D]) :-
    alignMapping1(Q, X),
    daughtersMapping1(Q, D).

synthOrInheritMap(T, M, U, N) :-    % Synthesised mapping.
    ground(T), ground(M), !,
    parentMap(T, M, U, N).
synthOrInheritMap(T, M, T, M).      % Inheritted mapping.

localMap(Q, V, D, ran(T, I), N) :-  % Local reference mapping.
    mapIdxSubscripts(1, D, L),
    alignment(V, ran(T, I), M),
    compose(M, L, N),
    mapDims(Q, I, N).

parentMap(T, M, T, M).              % Try to propagate synthesized mapping up.
parentMap(_, _, _, _).              % Else leave daughter/node maps un-unified.

mapIdxSubscripts(_, [], _).
mapIdxSubscripts(I, [S | D], M) :-
    mapIdxSubscript(I, S, M),
    J is I + 1,
    mapIdxSubscripts(J, D, M).

mapIdxSubscript(I, idx(X, _), M) :- member(X -> I, M).
mapIdxSubscript(_, _, _).

mapDims(_, _, []).
mapDims(Q, L, [_ -> D | M]) :-   % Align to array dimension
    member(D, L),
    mapDims(Q, L, M),
    notMember(_ -> D, M).
mapDims(Q, L, [_ -> ? | M]) :-   % Collapse Q or fewer indices
    Q > 0,
    R is Q - 1,
    mapDims(R, L, M).

compose(_, [], []).
compose(L, [X -> Y | M], [X -> Z | N]) :-
    ground(Y), !,
    member(Y -> Z, L),
    compose(L, M, N).
compose(L, [X -> _ | M], [X -> _ | N]) :- compose(L, M, N).


%enumList(_, [], []).
%enumList(I, [_ | D], [I | L]) :- J is I + 1, enumList(J, D, L).

% Phase III
% ---------
%   Complete instantiation of mappings at `arr' nodes.

alignMapping2(Q, fun(_, _, D, _, _)) :- daughtersMapping2(Q, D).
alignMapping2(Q, arr(_, D, _, align(ran(V, L), M))) :-
    ground(V),
    mapDims(Q, L, M),
    daughtersMapping2(Q, D).
alignMapping2(Q, var(_, D, _, _)) :- daughtersMapping2(Q, D).
alignMapping2(_, idx(_, _)).

daughtersMapping2(_, []).
daughtersMapping2(Q, [X | D]) :-
    alignMapping2(Q, X),
    daughtersMapping2(Q, D).

%-----------------------------------------------------------------

% Miscellaneous ``set'' operations.

union([], X, X).
union([X | R], Y, Z) :- member(X, Y), !, union(R, Y, Z). 
union([X | R], Y, [X | Z]) :- union(R, Y, Z).

subtract(M, [], M).
subtract(M, [X | I], N) :- remove(X, M, P), subtract(P, I, N).

member(X, [X | _]).
member(X, [_ | Y]) :- member(X, Y).

notMember(_, []).
notMember(X, [X | _]) :- !, fail.
notMember(X, [_ | L]) :- notMember(X, L).

remove(_, [], []).
remove(X, [X | L], L)       :- !.
remove(X, [Y | L], [Y | M]) :- remove(X, L, M).

%-----------------------------------------------------------------

% Allow successively more ``undefined'' (presumably collapsed)
% dimension-mappings, until success.

alignMap(P) :- alignMapping(0, P).

alignMapping(Q, P) :-
    alignMapping1(Q, P),
    alignMapping2(Q, P).
alignMapping(Q, P) :-
    R is Q + 1,
    alignMapping(R, P).

% Run the whole mapping procedure, reading the program from a file.

run(Tree, Alignments, Q) :-
    [Alignments],
    !,
    go(Tree, Q).
run(_, Alignments, _) :- printString("Failed to load file "), write(Alignments).

go(Tree, Q) :-
    see(Tree),
    getNode(P),
    !,
    seen,
    nl,
    alignDomain(P, Q, _, _),
    alignMap(Q),
    nl,
    printAssocs(0, Q).
go(Tree, _) :- printString("Failed to parse file "), write(Tree), seen.

