function D = pColon(a, varargin)
%CODISTRIBUTED.PCOLON Shared implementation of codistributed/colon equivalent to A:B or A:D:B
%   C = CODISTRIBUTED.COLON(A,B)
%   C = CODISTRIBUTED.COLON(A,D,B)
%   C = COLON(A,B) where A or B are codistributed array
%   
%   See also COLON, CODISTRIBUTED.


%   Copyright 2009-2022 The MathWorks, Inc.

narginchk(2, 5);

distributedutil.CodistParser.verifyNonCodistributedInputs([{a}, varargin]);

try
    [a, d, b, codistr, allowCommunication] = iParseArgs(a, varargin{:});
    if allowCommunication
        % Don't display CODISTRIBUTED.COLON to users because they might arrive here
        % via the colon overloads on codistributor1d and 2dbc.
        distributedutil.CodistParser.verifyReplicatedInputArgs('colon', ...
                                                          {a, d, b, codistr});
    end
catch e
    throw(e);
end

[LP, codistr] = codistr.hColonImpl(a, d, b);
% We have already verified that a, d, and b are called collectively, so no
% further error checking is necessary.
D = codistributed.pDoBuildFromLocalPart(LP, codistr);

end % End of codistributed.pColon.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [a, d, b, codistr, allowCommunication] = iParseArgs(a, varargin)
% iParseArgs Input should be input arguments 2 to end of the input arguments to
% codistributed.colon.

% First, identify all uses of 'noCommunication' and specifications of
% codistributors at the end of the argument list.
argList = varargin;
[argList, allowCommunication] = distributedutil.CodistParser.extractCommFlag(argList);
[argList, codistr] = distributedutil.CodistParser.extractCodistributor(argList);

if isempty(argList)
    error(message('parallel:distributed:ColonNoEndpoint')); 
end

% The only remaining options are:
% {b}
% {d, b}
switch length(argList)
  case 1
    % argList must be {b}.
    d = 1;
    b = argList{1};
  case 2
    % argList must be {d, b} 
    d = argList{1};
    b = argList{2};
  otherwise
    % argList{3} and up were invalid.
    error(message('parallel:distributed:ColonInvalidArgs'));
end

if ~allowCommunication
    distributedutil.CodistParser.verifyNotCodistWithNoComm('colon', ...
                                                      {a, b, d});
end

a = distributedutil.CodistParser.gatherIfCodistributed(a);
argList = distributedutil.CodistParser.gatherElements(argList);

% Throw the same error as MATLAB if not 2D (>2D is not detected by
% the mathcore helper).
if ~ismatrix(a) || ~ismatrix(d) || ~ismatrix(b)
    error(message('MATLAB:colon:inputsMustBe2D'));
end

% Use in-memory to parse the inputs and do the error checking. Note
% in-memory checker requires each input to have a valid type - char is
% supported, string is not.
isValidColonArgs = @(x) isinteger(x) || isfloat(x) || islogical(x) ...
    || ischar(x);
if isValidColonArgs(a) && isValidColonArgs(d) && isValidColonArgs(b)
    args = parallel.internal.array.colonGetArgs(a, argList{:});
    % Unpack the arguments.
    a = args(1);
    d = args(2);
    b = args(3);
else
    % Throw codistributed specific error message as it specifies the
    % supported type for codistributed.
    error(message('parallel:distributed:ColonInvalidStartEnd'));
end

% We make slightly more stringent requirements than the in-memory colon as 
% we do not support char
if ischar(a) || ischar(b) || ischar(d)
    error(message('parallel:distributed:ColonInvalidStartEnd'));
end

end % End of iParseArgs.
