function data = pCatToEdgeOfWorkerGrid(codistr, data, dim)
%pCatToEdgeOfWorkerGrid Concatenate input data along dimension to edge of worker grid 
% out = pCatToEdgeOfWorkerGrid(codistr, data, dim) concatenates data along the dimension dim
% to the edge of the worker grid in the dimension dim.  dim must be 1, 2, or both.
%
% For example, with a worker grid of [2, 3] and 'row' orientation, the layout of
% the workers in the worker grid is:
%     | 1 | 2 | 3 |
%     | 4 | 5 | 6 |
% In that case, out = pCatToEdgeOfWorkerGrid(codistr, data, 1) returns
% On worker 1: cat(1, data_on_wkr1, data_on_wkr2, data_on_wkr3)
% On worker 4: cat(1, data_on_wkr4, data_on_wkr5, data_on_wkr5)
% On workers 2, 3, 5 and 6: []
%
% Similarly, with the worker grid above, out = pCatToEdgeOfWorkerGrid(codistr, data, 2) 
% returns
% On worker 1: cat(2, data_on_wkr1, data_on_wkr4)
% On worker 2: cat(2, data_on_wkr2, data_on_wkr5)
% On worker 3: cat(2, data_on_wkr3, data_on_wkr6)
% On worker 4, 5 and 6: []
%
% Finally, out = pCatToEdgeOfWorkerGrid(codistr, data, [1 2]) 
% returns the entire result on worker 1:
% On worker 1: [data_on_wkr1, data_on_wkr4;
%               data_on_wkr2, data_on_wkr5;
%               data_on_wkr3, data_on_wkr6]
% On workers 2, 3, 4, 5 and 6: []

%   Copyright 2009-2022 The MathWorks, Inc.

if any(dim <= 0 | dim > 2)
    error(message('parallel:codistributors:c2dbcProcessAlongDimInvalidDimension'));
end

if all(codistr.WorkerGrid(dim) == 1)
    % Nothing to do.  All workers are already on the edge of the grid in this
    % dimension.
    return;
end

% We may need to combine rows, columns, or both
if any(dim == 1)
    proc = codistr.pSpmdIndexToProcessorCol(spmdIndex);
    data = iProcessOne(data, 1, proc, spmdIndex);
end
if any(dim == 2)
    proc = codistr.pSpmdIndexToProcessorRow(spmdIndex);
    data = iProcessOne(data, 2, proc, spmdIndex);
end

end

function data = iProcessOne(data, dim, proc, workerIdx)
% Process one concatenation
% Split communicators according to the proc values. Looking at our example
% above for dim = 1, the groups would be for workers 1:3 and 4:6
splitter = distributedutil.CommSplitter(proc, workerIdx); %#ok<NASGU>
% spmdIndex is now 1 for the lowest value of original worker index in each group.
% In our example, workers 1 and 4 now both have spmdIndex 1, workers 2 and 5
% both have spmdIndex 2, and workers 3 and 6 now have spmdIndex 3.

targetWorker = 1;
data = spmdCat(data, dim, targetWorker);
splitter = []; %#ok<NASGU> Undo communicator split.
end