function localLinInd = hFindDiagElementsInLocalPart(codistr)
%hFindDiagElementsInLocalPart Return the local linear indices of the diagonal 
% elements of the local part for codistributor1d.
% Note that this is restricted to finding the local diagonal elements of a
% 2D matrix.

%   Copyright 2010-2022 The MathWorks, Inc.
 
globalSz = codistr.Cached.GlobalSize;
par = codistr.Partition;
[e,f] = codistr.globalIndices(codistr.Dimension, spmdIndex);
cumPar = cumsum(par);
switch codistr.Dimension
    case 1
        offset = 0;
        if spmdIndex > 1
            offset = par(spmdIndex)*cumPar(spmdIndex-1);
        end
        localSz = codistr.hLocalSize();
        localLinInd = offset+1 : par(spmdIndex)+1 : min(prod(localSz),par(spmdIndex)^2+offset);
    case 2
        if e <= globalSz(1)
            if f <= globalSz(1)
                localLinInd = e:globalSz(1)+1:globalSz(1)*par(spmdIndex);
            else
                if (spmdIndex > 1)
                    offset = cumPar(spmdIndex-1);
                else
                    offset = 0;
                end
                localLinInd = e:globalSz(1)+1:globalSz(1)*(globalSz(1)-offset);
            end
        else
            localLinInd = [];
        end
    otherwise
        if par(spmdIndex) ~= 0
            localLinInd = 1:globalSz(1)+1:prod(globalSz);
        else
            localLinInd = [];
        end
end
