function [localV, szLeftPad, localX] = stencilCommPad(V, padSz, localX)
% STENCILCOMMPAD Adds padding to local parts of a distributed array V,
% returning the padded array and the number of elements added at its
% beginning (left end). If the local part of a second distributed array X
% is provided, that array is padded and returned also. V can be an array
% of any dimensionality distributed 1-d. X must be a column vector
% distributed with the same partition as V. The communication pattern is 
% optimized to avoid forwarding data. The algorithm used here must be kept
% in-sync with the similar algorithm in stencilCommConv.

%   Copyright 2022 The MathWorks, Inc.

hasX = nargin>2;
ourParts = cell(2, spmdSize);
needToRecv = false(1, spmdSize);
localV = getLocalPart(V);
codV = getCodistributor(V);
partitionV = codV.Partition;
cumPartitionV = [0 cumsum(partitionV)];
szVinDim = cumPartitionV(spmdSize+1);
partToSend = cell(2, 1);
dim = codV.Dimension;

% Figure out who is sending to us.
offset = cumPartitionV(spmdIndex);
ixLeftPad = max(offset+1-padSz(1), 1);  % Global index of first element in padded localX.
ixRightPad = min(offset+partitionV(spmdIndex)+padSz(2), szVinDim);  % Global index of last element in padded localX.
for spmdIx = 1:spmdIndex-1  % Loop over workers to our left.
    needToRecv(spmdIx) = cumPartitionV(spmdIx+1) >= ixLeftPad && partitionV(spmdIx);
end
for spmdIx = spmdIndex+1:spmdSize  % Loop over workers to our right.
    needToRecv(spmdIx) = cumPartitionV(spmdIx) < ixRightPad && partitionV(spmdIx);
end

% Do the sends and receives.
mwTag = 2^15 - mod(sum(mfilename),2^15);
for spmdIx = 1:spmdSize
    if spmdIx < spmdIndex
        padOffset = cumPartitionV(spmdIx+1) - offset;
        startIndex = 1;
        endIndex = min(padOffset + padSz(2), partitionV(spmdIndex));
    elseif spmdIx == spmdIndex
        continue  % Don't need to send anything to ourself.
    else % spmdIx > spmdIndex
        padOffset = cumPartitionV(spmdIx) + 1 - cumPartitionV(spmdIndex+1);
        startIndex = max(padOffset - padSz(1) + partitionV(spmdIndex), 1);
        endIndex = partitionV(spmdIndex);
    end
    % Communication happens here.
    needToSend = endIndex>=startIndex;
    if needToSend
        partToSend{1} = extractSlices(localV, startIndex:endIndex, dim);
        if hasX
            partToSend{2} = localX(startIndex:endIndex);
        end
        if needToRecv(spmdIx)
            ourParts(:, spmdIx) = spmdSendReceive(spmdIx, spmdIx, partToSend, mwTag+spmdIndex+spmdIx+spmdSize);
        else
            spmdSend(partToSend, spmdIx, mwTag+spmdIndex);
        end
    elseif needToRecv(spmdIx)
        ourParts(:, spmdIx) = spmdReceive(spmdIx, mwTag+spmdIx);
    else
        % Here if nothing to send or receive
    end
end
szUnPaddedLocalV = size(localV, dim);
for spmdIx = spmdIndex-1:-1:1  % Apply padding to left side of localV.
    localV = cat(dim, ourParts{1, spmdIx}, localV);
end
szLeftPad = size(localV, dim) - szUnPaddedLocalV;
for spmdIx = spmdIndex+1:spmdSize  % Apply padding to right side of localV.
    localV = cat(dim, localV, ourParts{1, spmdIx});
end
if hasX
    leftPadX = cat(1, ourParts{2, 1:spmdIndex-1});
    rightPadX = cat(1, ourParts{2, spmdIndex+1:spmdSize});
    localX = [leftPadX; localX; rightPadX];
end
end
