function C = convCommon(A, B, convFcnHandle, outClass, isSame, isValid, forceResultRow)
% Perform convolution of a distributed array A and another array B that may
% or may not be distributed.  forcesResultRow is an optional parameter used by
% conv (1D convolution) to return a result with the right shape.

% Copyright 2022 The MathWorks, Inc.

if nargin < 7
    forceResultRow = false;
end

if ~iscodistributed(A)
    [A, B] = deal(B, A);
    ABswapped = true;
else
    ABswapped = false;
end

assert(iscodistributed(A), 'convCommon should not be called if neither A nor B is codistributed')
codistrA = getCodistributor(A);
partA = codistrA.Partition;
cumPartA = [0 cumsum(partA)];
dim = codistrA.Dimension;
if ~isscalar(A) && isvector(A) && cumPartA(end) == 1  
    % A is a vector distributed in its singleton dimension.  Its transpose
    % will be used for calculations, so change its properties here.
    cumPartA = cumPartA*numel(A);
    dim = 3 - dim;
end
isConv1d = isequal(convFcnHandle, @conv);
if isConv1d
    dim = forceResultRow+1;
    szA = [1 1];
    szA(dim) = numel(A);
    szB = [1 1];
    szB(dim) = numel(B);
    sizeA = numel(A);
    sizeB = numel(B)-1;
else
    szA = size(A);
    sizeA = size(A, dim);
    if isstruct(B)
        szB = [length(B.hrow), length(B.hcol)];
    else
        szB = size(B);
    end
    szB = [szB ones(1, dim - length(szB))];
    sizeB = 0;
    if dim<=ndims(B)
        sizeB = max(szB(dim)-1, 0);
    end
end

isCodB = iscodistributed(B);
if isCodB
    temp = convDistDist(A, B, convFcnHandle, outClass, forceResultRow);
else
    temp = convDistRepl(A, B, convFcnHandle, outClass, forceResultRow);
end
szC = [size(temp) ones(1, dim - ndims(temp))];
szC(dim) = sizeA + sizeB;
if ~isCodB
    cumPartC = cumPartA;
    lastA = find(partA, 1, 'last') + 1;
    cumPartC(lastA:end) = szC(dim);
    partC = diff(cumPartC);
else
    partC = codistributor1d.defaultPartition(szC(dim));
end
cumPartC = [0 cumsum(partC)];
if isSame || isValid
    [cumPartC, temp, szC] = convTrimResult(cumPartC, temp, dim, szA, szB, szC,...
        isSame, isValid, ABswapped);
end
partC = diff(cumPartC);
codistrC = codistributor1d(dim, partC, szC);

% At this point, each worker has the convolution of the local part of A
% with all of B, whether it was distributed or not.  Now combine the
% results across workers, putting them on the worker where they belong.
ourParts = stencilCommConv(dim, temp, cumPartC, cumPartA, szB, isCodB);
offset = cumPartA(spmdIndex);
szLocalC = szC;
szLocalC(dim) = partC(spmdIndex);
szCInDim = szLocalC(dim);
ourPart = ourParts{spmdIndex};
if ~isempty(ourPart) && size(ourPart, dim) == szCInDim
    localC = ourPart;
else
    localC = zeros(szLocalC, 'like', temp);
    if ~isempty(ourPart)
        if cumPartC(spmdIndex) < offset
            localC = assignSlices(localC, 1:size(ourPart, dim), ourPart, dim);
        else
            localC = assignSlices(localC, (szCInDim-size(ourPart, dim)+1):szCInDim, ourPart, dim);
        end
    end
end

for ii = 1:spmdSize
    if ii == spmdIndex
        continue
    end
    toAdd = ourParts{ii};
    if ~isempty(toAdd)
        if size(toAdd, dim) == szCInDim
            localC = localC + toAdd;
        elseif size(toAdd, dim) < szCInDim
            if ii <= spmdIndex  % Came from left.  Add at beginning.
                changedSlices = 1:size(toAdd, dim);
            else % Came from right, Add at end.
                changedSlices = (szCInDim - size(toAdd, dim) + 1):szCInDim;
            end
            if ~isempty(localC)
                priorC = extractSlices(localC, changedSlices, dim);
                localC = assignSlices(localC, changedSlices, priorC + toAdd, dim);
            else
                localC = toAdd;
            end
        end
    end
end
C = codistributed.pDoBuildFromLocalPart(localC, codistrC);
end
