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

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

% Index Mapping
% -------------

% Glossary of symbols...
%
% V       = variable
% F       = function
% X, Y, Z = index
% T       = template
% I       = template dimension
% J       = variable dimension
%
% D, E           = (daughter) node---complicated syntax!
%
% K, L, M, N     = index to template dimension mappings---[] | [X -> I | M]
% P, Q, R, S     = template dimension to index mappings---[] | [I -> X | P]
% A, B           = full index map---map(T, M)
 

% Preprocessing
% -------------
%      Preprocess expression tree prior to mapping.  Turn `fun' and `var'
% nodes into leaves.  After processing, the only true non-leaf node is
% `nod'.  `arg' is regarded as a modification of a `nod' daughter.
% `arc' is regarded as a modifier on an arc.
%
% Analyse subscript expressions for linearity.
% Set up alignment maps for arguments.

preprocess(var(V, Ds), arc(nod([arc(var(V, Subs)) | Es]))) :-
    preprocessSubs(Ds, Subs, Es).
preprocess(fun(F, Xs, Ds), arc(nod([fun(F, Xs, M), nod(Es)]))) :-
    function(F, Aligns),
    resolveAligns(Aligns, [variable(_, _, Js) | Resolves]),
    map(Xs, Js, M),
    preprocessArgs(Ds, Es, Resolves).
preprocess(idx(X), idx(X)).

preprocessSubs([], [], []).
preprocessSubs([D | Ds], [Sub | Subs], [E | Es]) :-
    analyseSub(D, Sub),
    preprocess(D, E),
    preprocessSubs(Ds, Subs, Es).

preprocessArgs([], [], []).
preprocessArgs([arg(Xs, D) | Ds], [arg(Xs, M, E) | Es],
               [variable(_, _, Js) | Resolves]) :-
    map(Xs, Js, M),
    preprocess(D, E),
    preprocessArgs(Ds, Es, Resolves).


% Analyse subscript expression for linearity
% (Will eventually support more general definition of ``linearity'').

analyseSub(idx(X), lin(X)) :- !.
analyseSub(_, nonlin).


% Resolve a list of (argument) alignments.  Eg
%
%  [variable(res, T,    _),
%   variable(a,   res, [1, *]),
%   variable(b,   res, [*, 2])]
%
% yields
%
%  [variable(res, T, [I, J | _]),
%   variable(a,   T, [I, *]),
%   variable(b,   T, [*, J])].
%
% The alignment chain is terminated on targets which (are uninstantiated
% or) don't appear in the original list.

resolveAligns(Aligns, Resolves) :-
    sameKeys(Aligns, Resolves),
    composeAligns(Resolves, Aligns, Resolves).


sameKeys([], []).
sameKeys([variable(V, _, _) | As], [variable(V, _, _) | Bs]) :- sameKeys(As, Bs).


composeAligns(_, [], []).
composeAligns(Resolves, [Align | AlignsTail], [Resolve | ResolvesTail]) :-
    composeAlign(Resolves, Align, Resolve),
    composeAligns(Resolves, AlignsTail, ResolvesTail).

composeAlign(Resolves, variable(V, T, Is), variable(V, U, Ks)) :-
    ground(T),
    member(variable(T, U, Js), Resolves), !,
    compose(Is, Js, Ks).
composeAlign(_, variable(_, T, Is), variable(_, T, Is)).


% First argument of `compose' should be fully instantiated.  Instantiation
% status of other arguments is arbitrary.

compose([], _, []).
compose([I | Is], Js, [K | Ks]) :- isNth(I, Js, K), compose(Is, Js, Ks).

isNth(*, _, *)        :- !.
isNth(1, [X | _ ], X) :- !.
isNth(N, [_ | Xs], X) :- N1 is N - 1, isNth(N1, Xs, X).

append([], Append, Append).
append([X | Old], Append, [X | New]) :- append(Old, Append, New).

% Decorate
% --------
%      Decorate tree with uninstantiated mapping attribute fields.

decorate(nod(Ds), nod(Es, _, _, _)) :-
    decorateDaughters(Ds, Es).
decorate(arg(Xs, M, D), arg(Xs, M, E, _, _, _)) :-
    decorate(D, E).
decorate(arc(D), arc(E, _, _, _)) :-
    decorate(D, E).
decorate(var(V, Subs), var(V, Subs, _, _)).
decorate(fun(F, Xs, M), fun(F, Xs, M)).
decorate(idx(X), idx(X)).

decorateDaughters([], []).
decorateDaughters([D | Ds], [E | Es]) :-
    decorate(D, E),
    decorateDaughters(Ds, Es).


% Postprocessing
% --------------
%      Reverse preprocessing transformations on the decorated syntax tree,
% ready for printing.

postprocess(arc(nod([arc(var(V, _, _, _), _, _, _) | Ds], _, A, _), _, _, _),
            var(V, Es, A)) :-
    postprocessDaughters(Ds, Es).
postprocess(arc(nod([fun(F, Xs, _), nod(Ds, _, _, _)], _, A, _), _, _, _),
            fun(F, Xs, Es, A)) :-
    postprocessDaughters(Ds, Es).
postprocess(arg(Xs, _, D, _, A, _), arg(Xs, E, A)) :-
    postprocess(D, E).
postprocess(idx(X), idx(X)).

postprocessDaughters([], []).
postprocessDaughters([D | Ds], [E | Es]) :-
    postprocess(D, E),
    postprocessDaughters(Ds, Es).

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

% Index mapping
% -------------

% Phase I
% -------
%      Instantiate map domains at each node.


% Phase Ia: domains of `nod' maps constructed from children.
% Post-order traversal.

idxMapDomainA(nod(Ds, Xs, _, _), Xs)   :- idxMapDomainADaughters(Ds, Xs).
idxMapDomainA(arg(Xs, _, D, _, _, _), Zs) :-
    idxMapDomainA(D, Ys),
    subtract(Ys, Xs, Zs).
idxMapDomainA(arc(D, _, _, _), Xs)     :- idxMapDomainA(D, Xs).
idxMapDomainA(var(_, _, _, _), []).
idxMapDomainA(fun(_, Xs, _), Xs).
idxMapDomainA(idx(X), [X]).
    
idxMapDomainADaughters([], []).
idxMapDomainADaughters([D | Ds], Zs) :-
    idxMapDomainA(D, Xs),
    idxMapDomainADaughters(Ds, Ys),
    union(Xs, Ys, Zs).


% Phase Ib: domains of `arg' and `var' maps constructed from parent.
% Pre-order traversal.

idxMapDomainB(nod(Ds, Xs, _, _), _)     :- idxMapDomainBDaughters(Ds, Xs).
idxMapDomainB(arg(Xs, _, D, Zs, _, _), Ys) :-
    union(Xs, Ys, Zs),
    idxMapDomainB(D, Zs).
idxMapDomainB(arc(D, _, _, _), Xs)      :- idxMapDomainB(D, Xs).
idxMapDomainB(var(_, _, Xs, _), Xs).
idxMapDomainB(fun(_, _, _), _).
idxMapDomainB(idx(_), _).
    
idxMapDomainBDaughters([], _).
idxMapDomainBDaughters([D | Ds], Xs) :-
    idxMapDomainB(D, Xs),
    idxMapDomainBDaughters(Ds, Xs).


% Phase Ic: domains of `arc' maps constructed from children.
% Post-order traversal.

idxMapDomainC(nod(Ds, Xs, _, _), Xs)   :- idxMapDomainCDaughters(Ds).
idxMapDomainC(arg(_, _, D, Xs, _, _), Xs) :- idxMapDomainC(D, _).
idxMapDomainC(arc(D, Xs, _, _), Xs)    :- idxMapDomainC(D, Xs).
idxMapDomainC(var(_, _, Xs, _), Xs).
idxMapDomainC(fun(_, Xs, _), Xs).
idxMapDomainC(idx(X), [X]).
    
idxMapDomainCDaughters([]).
idxMapDomainCDaughters([D | Ds]) :-
    idxMapDomainC(D, _),
    idxMapDomainCDaughters(Ds).


% Phase II
% --------
%      Impose horizontal constraints on maps.  Pre-order traversal.


idxMapConstraint(nod(Ds, Xs, map(T, M), _), T, N)        :-
    domain(M, Xs),
    subset(M, N),
    idxMapConstraintDaughters(Ds, T, M).
idxMapConstraint(arg(_, L, D, Xs, map(T, M), _), T, N)   :-
    domain(M, Xs),
    subset(L, M),
    subset(N, M),
    idxMapConstraint(D, T, M).
idxMapConstraint(arc(D, Xs, map(T, M), map(U, N)), U, L) :-
    domain(N, Xs),
    subset(N, L),
    domain(M, Xs),
    idxMapConstraint(D, T, M).
idxMapConstraint(var(_, _, Xs, map(T, M)), T, N)         :-
    domain(M, Xs),
    subset(M, N).
idxMapConstraint(fun(_, _, L), _, M) :-
    subset(L, M).
idxMapConstraint(idx(_), _, _).
    
idxMapConstraintDaughters([], _, _).
idxMapConstraintDaughters([D | Ds], T, M) :-
    idxMapConstraint(D, T, M),
    idxMapConstraintDaughters(Ds, T, M).


% Phase III
% ---------
%      Index mapping proper.


% Phase IIIpre: Preamble.  Insert mapping ``quality'' in `nod' and `arg' nodes.
% Currently just the maximal template rank.

idxMapQual(D) :-
    idxMapMaxDim(D, Max),
    idxMapQualNode(Max, D).


idxMapMaxDim(nod(Ds, _, _, _), Max)      :- idxMapMaxDimDaughters(Ds, Max).
idxMapMaxDim(arg(_, _, D, _, _, _), Max) :- idxMapMaxDim(D, Max).
idxMapMaxDim(arc(D, _, _, _), Max)       :- idxMapMaxDim(D, Max).
idxMapMaxDim(var(V, _, _, _), Max)       :- templateDim(V, Max).
idxMapMaxDim(fun(_, _, _), 0).
idxMapMaxDim(idx(_), 0).

idxMapMaxDimDaughters([], 0).
idxMapMaxDimDaughters([D | Ds], Max) :-
    idxMapMaxDim(D, Max1),
    idxMapMaxDimDaughters(Ds, Max2),
    Max is max(Max1, Max2).


templateDim(V, Dim) :-
    alignment(V, T, _),
    variable(T, _, Is),
    size(Is, Dim).


idxMapQualNode(Max, nod(Ds, _, _, Max))   :-
     idxMapQualNodeDaughters(Max, Ds).
idxMapQualNode(Max, arg(_, _, D, _, _, Max)) :-
     idxMapQualNode(Max, D).
idxMapQualNode(Max, arc(D, _, _, _)) :- idxMapQualNode(Max, D).
idxMapQualNode(_, var(_, _, _, _)).
idxMapQualNode(_, fun(_, _, _)).
idxMapQualNode(_, idx(_)).

idxMapQualNodeDaughters(_, []).
idxMapQualNodeDaughters(Max, [D | Ds]) :-
    idxMapQualNode(Max, D),
    idxMapQualNodeDaughters(Max, Ds).


% Phase IIIa: Apply `var' heuristic.

idxMapVar(nod(Ds, _, _, _))      :- idxMapVarDaughters(Ds).
idxMapVar(arg(_, _, D, _, _, _)) :- idxMapVar(D).
idxMapVar(arc(D, _, _, _))       :- idxMapVar(D).
idxMapVar(var(V, Subs, _, A))    :- varSynthMap(V, Subs, A).
idxMapVar(fun(_, _, _)).
idxMapVar(idx(_)).

idxMapVarDaughters([]).
idxMapVarDaughters([D | Ds]) :-
    idxMapVar(D),
    idxMapVarDaughters(Ds).


varSynthMap(V, Subs, map(T, M)) :-
    alignment(V, T, Is),
    mapLinSubscripts(Subs, Is, M).

mapLinSubscripts([], [], _).
mapLinSubscripts([Sub | Subs], [I | Is], M) :-
    mapLinSubscript(Sub, I, M),
    mapLinSubscripts(Subs, Is, M).

mapLinSubscript(lin(X), I, M) :- member(X -> I, M).
mapLinSubscript(_, _, _).


alignment(V, U, Ks) :-
    variable(V, T, Is),
    ground(T), !,
    alignment(T, U, Js),
    compose(Is, Js, Ks).
alignment(V, V, Is) :- variable(V, _, Is), enumFrom(1, Is).

enumFrom(_, []).
enumFrom(N, [N | Is]) :- N1 is N + 1, enumFrom(N1, Is).


% Phase IIIb: Apply `arc' heuristic.

idxMapArc(nod(Ds, _, _, _))      :- idxMapArcDaughters(Ds).
idxMapArc(arg(_, _, D, _, _, _)) :- idxMapArc(D).
idxMapArc(arc(D, _, A, B))       :- idxMapArc(D), arcSynthMap(A, B).
idxMapArc(var(_, _, _, _)).
idxMapArc(fun(_, _, _)).
idxMapArc(idx(_)).

idxMapArcDaughters([]).
idxMapArcDaughters([D | Ds]) :-
    idxMapArc(D),
    idxMapArcDaughters(Ds).


arcSynthMap(A, A).
arcSynthMap(_, _).


% Phase IIIc: Instantiate remaining mappings.

% Note that with this strategy, it's known that some expressions
% can't be mapped at all!  This needs fixing eventually.

idxMapAll(nod(Ds, _, A, Qual))   :-
    nodSynthMap(Qual, A),
    idxMapAllDaughters(Ds).
idxMapAll(arg(_, _, D, _, A, Qual)) :-
    nodSynthMap(Qual, A),
    idxMapAll(D).
idxMapAll(arc(D, _, _, _)) :- idxMapAll(D).
idxMapAll(var(_, _, _, _)).
idxMapAll(fun(_, _, _)).
idxMapAll(idx(_)).

idxMapAllDaughters([]).
idxMapAllDaughters([D | Ds]) :-
    idxMapAll(D),
    idxMapAllDaughters(Ds).


nodSynthMap(Qual, map(T, M)) :-
    ground(T),
    alignment(T, T, Is),
    mapAll(Qual, Is, M).

mapAll(Max, Is, M) :-
    size(Is, TargDim),
    countUncollapsed(M, NumUncollapsed),
    NumFailures is max(0, NumUncollapsed - TargDim),
    (NumFailures =:= 0, ! | TargDim =:= Max),
    mapIndices(Is, NumFailures, M).


mapIndices(_, 0, []).
mapIndices(Is, N, [_ -> I | M]) :-
    mapIndex(Is, N, Is1, N1, I),
    mapIndices(Is1, N1, M).

mapIndex(Is, N, Is,  N,  I) :- ground(I), I = *, !.
mapIndex(Is, N, Is1, N,  I) :- removeMember(I, Is, Is1).
mapIndex(Is, N, Is,  N1, ?) :- N > 0, N1 is N - 1.


countUncollapsed([], 0).
countUncollapsed([_ -> I | M], N)  :-
    ground(I), I = *, !,
    countUncollapsed(M, N).
countUncollapsed([_ -> _ | M], N1) :- countUncollapsed(M, N), N1 is N + 1.


% nod(Ds, Xs, A, Qual)
% arg(Xs, M, D, Ys, A, Qual)
% arc(D, Xs, A, B)
% var(V, Subs, Xs, A)
% fun(F, Xs, M)
% idx(X)

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

% Embedding
% ---------

% Phase I
% -------
%      Decorate tree with ``embedding'' records.

embedDecorate(nod(Ds, Xs, A, _), nod(Es, B)) :-
    invertIdxMap(Xs, A, B),
    embedDecorateDaughters(Ds, Es).
embedDecorate(arg(Xs, M, D, Ys, A, _), arg(Xs, M, E, B)) :-
    invertIdxMap(Ys, A, B),
    embedDecorate(D, E).
embedDecorate(arc(D, Xs, A1, A2), arc(E, B1, B2)) :-
    invertIdxMap(Xs, A1, B1),
    invertIdxMap(Xs, A2, B2),
    embedDecorate(D, E).
embedDecorate(var(V, Subs, Xs, A), var(V, Subs, B)) :-
    invertIdxMap(Xs, A, B).
embedDecorate(fun(F, Xs, M), fun(F, Xs, M)).
embedDecorate(idx(X), idx(X)).
 
embedDecorateDaughters([], []).
embedDecorateDaughters([D | Ds], [E | Es]) :-
    embedDecorate(D, E),
    embedDecorateDaughters(Ds, Es).


invertIdxMap(Xs, map(T, M), embed(Xs, T, P)) :-
    variable(T, _, P),
    invert(M, P).

invert([], _).
invert([X -> I | M], P) :- invertMapping(X, I, P), invert(M, P).

invertMapping(_, *, _) :- !.
invertMapping(_, ?, _) :- !.
invertMapping(X, I, P) :- isNth(I, P, X).


embedPostprocess(arc(nod([arc(var(V, _, _), _, _) | Ds], A), _, _),
		 var(V, Es, A)) :-
    embedPostprocessDaughters(Ds, Es).
embedPostprocess(arc(nod([fun(F, Xs, _), nod(Ds, _)], A), _, _),
                 fun(F, Xs, Es, A)) :-
    embedPostprocessDaughters(Ds, Es).
embedPostprocess(arg(Xs, _, D, A), arg(Xs, E, A)) :-
    embedPostprocess(D, E).
embedPostprocess(idx(X), idx(X)).

embedPostprocessDaughters([], []).
embedPostprocessDaughters([D | Ds], [E | Es]) :-
    embedPostprocess(D, E),
    embedPostprocessDaughters(Ds, Es).

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

% 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).

intersection(X, Y, Z) :- subtract(X, Y, W), subtract(X, W, 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).

removeMember(X, [X | L], L).
removeMember(X, [Y | L], [Y | M]) :- removeMember(X, L, M).

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).

subset([], _).
subset([X | M], N) :- member(X, N), !, subset(M, N).

size([], 0).
size([_ | R], N) :- ground(N), !, N > 0, N1 is N - 1, size(R, N1).
size([_ | R], N) :- size(R, N1), N is N1 + 1.

map([], [], []).
map([X | Xs], [Y | Ys], [X -> Y | M]) :- map(Xs, Ys, M).

domain([], []).
domain([X -> _ | M], [X | Xs]) :- domain(M, Xs).

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

% 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, R1) :-
%go(Tree, R) :-
%go(Tree, Q1) :-
%go(Tree, P1) :-
%go(Tree, Q) :-
%go(Tree, P2) :-
    see(Tree),
    getNode(P),
    !,
    seen,
    nl,
    preprocess(P, P1),
    decorate(P1, P2),
    idxMapDomainA(P2, _),
    idxMapDomainB(P2, _),
    idxMapDomainC(P2, _),
    idxMapConstraint(P2, _, _),
    idxMapQual(P2),
    idxMapVar(P2),
    idxMapArc(P2),
    idxMapAll(P2),
    embedDecorate(P2, R),

%    nl,
%    printNodArc(0, P2).

%    postprocess(P2, Q),
%    nl,
%    printIdxMap(0, Q).

    embedPostprocess(R, R1),
    nl,
    printEmbed(0, R1).
go(Tree, _) :- printString("Failed to parse file "), write(Tree), seen.

