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

% Index Dependency Analysis


depen(arr(Is, E, Atr), Xs) :-
  depen(E, Ys),
  union(Ys, Is, Zs),       % Make sure all locally bound indices included
  atrISet(Atr, Zs),
  subtract(Zs, Is, Xs).    % Locally bound indices not passed up to parent
depen(idx(X), [X]).
depen(var(_, Es, Atr), Xs) :-
  depenList(Es, Xss),
  atrSubsISets(Atr, Xss),  % Save for future reference
  unionList(Xss, Xs),
  atrISet(Atr, Xs).
depen(fun(_, Es, Atr), Xs) :-
  depenList(Es, Xss),
  unionList(Xss, Xs),
  atrISet(Atr, Xs).

depenList([], []).
depenList([E | Es], [Xs | Xss]) :-
  depen(E, Xs),
  depenList(Es, Xss).
 

% Horizontal Constraints


constrain(arr(Is, E, Atr), T, R, P) :-  % T: variable, rank R.  P: mappings
  atrISet(Atr, Xs), subtract(Xs, Is, Ys),
				% `Ys': idxs *not* bound by this arr node
  domain(L, Ys), subset(L, P),  % `L': map with domain `Ys', contained in `P'
  domain(M, Xs), subset(L, M),  % `M': map with domain `Xs', containing `L'
  atrNodMap(Atr, map(T, R, M)),
  constrain(E, T, R, M).
constrain(idx(_), _, _, _).
constrain(var(_, Es, Atr), T, R, P) :-
  atrISet(Atr, Xs),
  domain(L, Xs), subset(L, P),  % `L': map with domain `Xs', contained in `P'
  atrDstMap(Atr, map(T, R, L)),
  domain(M, Xs),
  atrSrcMap(Atr, map(U, E, M)), % `M': map with domain `Xs'
  constrainList(Es, U, E, M).
constrain(fun(_, Es, Atr), T, R, P) :-
  atrISet(Atr, Xs),
  domain(L, Xs), subset(L, P),  % `L': map with domain `Xs', contained in `P'
  atrDstMap(Atr, map(T, R, L)),
  domain(M, Xs),
  atrSrcMap(Atr, map(U, E, M)), % `M': map with domain `Xs'
  constrainList(Es, U, E, M).

constrainList([], _, _, _).
constrainList([E | Es], T, R, P) :-
  constrain(E, T, R, P),
  constrainList(Es, T, R, P).


% Variable Reference Mapping.


varMap(arr(_, E, _))    :- varMap(E).
varMap(idx(_)).
varMap(var(V, Es, Atr)) :-
  % Attempt ``local reference mapping''...
  atrSubsISets(Atr, Xss),
  size(Es, R),          % Compute rank of `V' from number of subscripts
  atrSrcMap(Atr, map(V, R, M)),
  mapSingletSubscripts(Xss, M, 1),
  varMapList(Es).
varMap(var(_, Es, _))   :-
  % ...or leave uninstantiated.
  varMapList(Es).
varMap(fun(_, Es, _))   :-  varMapList(Es).

varMapList([]).
varMapList([E | Es]) :- varMap(E), varMapList(Es).

mapSingletSubscripts([], _, _).
mapSingletSubscripts([Xs | Xss], M, D) :-
  mapSingletSubscript(Xs, M, D),
  D1 is D + 1, mapSingletSubscripts(Xss, M, D1).

mapSingletSubscript([X], M, D) :-
  % Attempt to map index `X' to dimension `D'...
  member(X -> D, M).
mapSingletSubscript(_, _, _).
  % ...or leave mapping to `D' uninstantiated.


% Propagation of Mappings.


propagate(arr(_, E, _))    :- propagate(E).
propagate(idx(_)).
propagate(var(_, Es, Atr)) :-
  propagateList(Es),
  %atrDstMap(Atr, map(T, R, L)),
  %atrSrcMap(Atr, map(U, E, M)),
  %propagateArc(T, R, L, U, E, M).
  atrDstMap(Atr, MD),
  atrSrcMap(Atr, MS),
  propagateArc(MD, MS).
propagate(fun(_, Es, Atr))   :-
  propagateList(Es),
  %atrDstMap(Atr, map(T, R, L)),
  %atrSrcMap(Atr, map(U, E, M)),
  %propagateArc(T, R, L, U, E, M).
  atrDstMap(Atr, MD),
  atrSrcMap(Atr, MS),
  propagateArc(MD, MS).

propagateList([]).
propagateList([E | Es]) :- propagate(E), propagateList(Es).

%propagateArc(T, R, M, T, R, M).  % propagate mapping...
%propagateArc(_, _, _, _, _, _).  % ...or don't

propagateArc(MD, MD).  % propagate mapping...
propagateArc(_, _).    % ...or don't


% Range Injection.


select(arr(_, E, Atr), CRank) :-
  atrNodMap(Atr, map(T, R, M)),
  ground(T),
  mapAll(M, R, CRank),
  select(E, CRank).
select(idx(_), _).
select(var(_, Es, Atr), CRank) :-
  atrSrcMap(Atr, map(T, R, M)),
  ground(T),
  mapAll(M, R, CRank),
  selectList(Es, CRank).
select(fun(_, Es, Atr), CRank) :-
  atrSrcMap(Atr, map(T, R, M)),
  ground(T),
  mapAll(M, R, CRank),
  selectList(Es, CRank).

selectList([], _).
selectList([E | Es], CRank) :- select(E, CRank), selectList(Es, CRank).

mapAll(M, R, CRank) :-
  R < CRank, !,
  enum(R, Dims),        % enum(N, [1, ..., N]).
  inject(M, Dims).
mapAll(M, R, _) :-
  enum(R, Dims),
  injectOrCollapse(M, Dims).

inject([], _).
inject([_ -> D | M], Range) :-
  member(D, Range),
  remove(D, Range, Range1),
  inject(M, Range1).

injectOrCollapse([], _).
injectOrCollapse([_ -> D | M], Range) :-
  member(D, Range),
  remove(D, Range, Range1),
  injectOrCollapse(M, Range1).
injectOrCollapse([_ -> * | M], Range) :-
  injectOrCollapse(M, Range).


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

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

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.

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

unionList([], []).
unionList([S | Ss], U) :- unionList(Ss, V), union(S, V, U).

enum(N, Is) :- enumFrom(1, N, Is).

enumFrom(I, N, [I | Is]) :- I =< N, !, I1 is I + 1, enumFrom(I1, N, Is).
enumFrom(_, _, []).

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

% Selector Predicates.

atrISet(atr(Xs, _, _, _), Xs).
atrSubsISets(atr(_, Xss, _, _), Xss).
atrNodMap(atr(_, _, M, _), M).
atrSrcMap(atr(_, _, M, _), M).
atrDstMap(atr(_, _, _, M), M).

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

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

run(Tree, P, CRank) :-
    see(Tree),
    getNode(P),
    !,
    seen,
    nl,
    depen(P, _),
    constrain(P, _, _, _),
    varMap(P),
    propagate(P),
    select(P, CRank),
    printIdxMap(0, P),
    true.

