function [LP, codistr] = hReductionOpAlongDimImpl(codistr, fcn, LP, dim) 
%hReductionOpAlongDimImpl Implementation for codistributor2dbc.

%   Copyright 2009-2022 The MathWorks, Inc.

if isequal(dim, 0) && distributedutil.Sizes.isSquareEmptyMatrix(codistr.Cached.GlobalSize)
    % For 0-by-0 empty codistributed array, we have to call fcn with only 1 input
    % argument because neither all(zeros(0, 0), 1) nor all(zeros(0, 0), 2)
    % is equivalent to all(zeros(0, 0)). 
    % If the global size is 0-by-0 then all the local parts are also 0-by-0.
    LP = fcn(LP);
    % At this point, LP is replicated, so we convert it into the actual 
    % local part.
    codistr = codistributor2dbc();
    srcWorker = 0;  % All workers have the replicated array.
    [LP, codistr] = codistr.hBuildFromReplicatedImpl(srcWorker, LP);
    return
end

if isequal(dim, 0)
    dim = distributedutil.Sizes.firstNonSingletonDimension(codistr.Cached.GlobalSize);
end

% The reduction takes place in two stages.  Local reduction and a communicating
% reduction.  Look at a 2-by-2 worker grid to understand why:
% | 1 | 2 | 1 | 2 | 1 | 2 | ...
% | 3 | 4 | 3 | 4 | 3 | 4 | ...
% | 1 | 2 | 1 | 2 | 1 | 2 | ...
% | 3 | 4 | 3 | 4 | 3 | 4 | ...
% ...
%
% Assume the reduction is taking place along the rows.  By having all the
% workers reduce their local part along the reduction dimension, we end up
% with:
%
% | 1 | 2 | 1 | 2 | 1 | 2 | ...
% | 3 | 4 | 3 | 4 | 3 | 4 | ...

% 
% where each block is of size 1-by-BlockSize. 

% If using "all", expand to a numeric vector now
if matlab.internal.datatypes.isScalarText(dim) && strncmpi(dim, "all", strlength(dim))
    dim = 1:numel(codistr.Cached.GlobalSize);
end

% Local reduction step.
LP = fcn(LP, dim);

% Resolve cases that can be done without global communication.
if all(dim > 2)
    % The reduction is complete and the global size hasn't changed, so
    % codistributor does not need to be modified.
    return;
end

% In order to deal with vectors of dimensions correctly, we need to specify
% which reduction dimensions are the real ones, meaning the dimensions that
% actually exist. For 2dbc, this is only 1 and 2.
realDim = dim(dim<3);
orgSizeInDim = codistr.Cached.GlobalSize(realDim);
gsize = codistr.Cached.GlobalSize;
% The resulting global size is 1 in the reduction dimension, unmodified in
% other dimensions.
gsize(realDim) = 1;
codistr = codistributor2dbc(codistr.WorkerGrid, codistr.BlockSize, ...
    codistr.Orientation, gsize);

if all(codistr.WorkerGrid(realDim) == 1)
    % Each row/column resides in its entirety on a single worker, and all
    % workers are on the edge of the grid.
    return;
elseif orgSizeInDim <= codistr.BlockSize
    LP = iHandleInteriorWorkers(codistr, LP, realDim);
else
    LP = iReduceAcrossWorkers(codistr, fcn, LP, realDim);
end

end % End of hReductionOpAlongDimImpl.

function LP = iReduceAcrossWorkers(codistr, fcn, LP, dim)
% The input LP and codistr are the local part and codistributor after reducing
% along dim.  Handle the non-trivial case of reducing along the rows or the
% columns.
templateElement = LP;

% Concatenate all the reduced local parts along first dimension in dim, and
% store on the workers that are on the edge of the grid in the dim(1)-th
% dimension.  In our example, this stores the results on workers 1 and 2.
LP = codistr.pCatToEdgeOfWorkerGrid(LP, dim(1));

% If we have dim = [1 2], perform this step again in the second dimension
if numel(dim)==2
    LP = codistr.pCatToEdgeOfWorkerGrid(LP, dim(2));
end

% If the worker is on the grid edge in direction of dim(1) (and
% automatically also in direction of dim(2) if it exists), perform the
% final reduction.
if codistr.pIsOnWorkerGridEdge(dim(1))
    LP = fcn(LP, dim);
else
    % We pushed all the data to the edge of the grid, leaving us with
    % an empty local part of the appropriate size.  
    LP = distributedutil.Allocator.create(codistr.hLocalSize(), templateElement);
end

end % End of iReduceAcrossWorkers

function LP = iHandleInteriorWorkers(codistr, LP, dim)
% The input LP and codistr are the local part and codistributor after reducing
% along dim.  Handle the special case where the local part was non-trivial only
% on the edge of the grid, so we don't need to communicate to complete the
% reduction.
if ~codistr.pIsOnWorkerGridEdge(dim(1))
    LP = distributedutil.Allocator.create(codistr.hLocalSize(), LP);
end

end % End of iHandleInteriorWorkers
