% ProcessRemoteResourceSet object.


% Copyright 2021-2022 The MathWorks, Inc.

classdef ProcessRemoteResourceSet < spmdlang.AbstractRemoteResourceSet

    methods (Access = public, Static)

        % This function exists mainly to support parallel pool jobs, where the effective
        % world isn't 'world'.
        function x = getWorldCommsCell(nWorkers)
            [wc, workerIdFirst] = spmdlang.commForWorld('get');

            % Need to pad the cell array so that the workers can index into it by their
            % "World" spmdIndex.
            padCells = cell(1, workerIdFirst - 1);

            if ischar(wc) % probably 'world'
                x = repmat({wc}, 1, nWorkers);
            else
                x = wc;
            end
            x = [padCells, x]; % in parallel pool job, this is e.g.: {[], 6, 6, 6}
        end

        % Bootstrap - this static method builds the "world" set from the pool, and
        % requires no communication or remote execution to achieve that.
        function worldSet = buildWorldSet(pool)
            if isempty(pool)
                error(message('parallel:lang:spmd:NoSession'));
            end
            ps = pool.NumWorkers;
            isWorld = true;
            worldSet = spmdlang.ProcessRemoteResourceSet(...
                ps, spmdlang.ProcessRemoteResourceSet.getWorldCommsCell(ps), pool, isWorld);
        end
    end

    methods (Access = public, Hidden)
        function obj = ProcessRemoteResourceSet(numWorkers, commCell, pool, isWorld)
            obj = obj@spmdlang.AbstractRemoteResourceSet(numWorkers, commCell, pool, isWorld);
        end

        function invalidate(obj, workersStillAvailable)
            if ~obj.IsWorld && workersStillAvailable
                % Use SPMD itself to release the communicators
                spmd_feval_fcn(@iFreeMyComm, {}, obj);
            end
            % Finally, call the superclass invalidation
            invalidate@spmdlang.AbstractResourceSet(obj, workersStillAvailable);
        end

        % "Factory method" to build a remote resource set of a specific size from the
        % open pool. It does that by splitting the base set.
        % "sz" is simply the required size of the subset
        function subSetObj = buildRemoteResourceSet(obj, sz)

            if sz == obj.NumWorkers
                subSetObj = obj;
            elseif sz < 0
                error(message('parallel:lang:spmd:InvalidResourceSelection', sz));
            elseif sz < obj.NumWorkers

                % Build from the base set,
                newComms = spmd_feval_fcn(@iBuildCommHandle, { sz }, obj);

                % Pad as necessary to handle mismatch between MPI 'world' spmdIndex and
                % processes. (Happens in parallel pool jobs)
                [~, workerIdFirst] = spmdlang.commForWorld('get');
                padCells = cell(1, workerIdFirst - 1);
                newComms = {padCells{:}, newComms{:}}; %#ok<CCAT> Dereference the Composite returned from the spmd_feval_fcn

                isWorld = false; % trivially as sz < obj.NumWorkers
                subSetObj = spmdlang.ProcessRemoteResourceSet(sz, newComms(:), obj.Pool, isWorld);
            else
                % Should never get here because there's already a check that will not let us
                % build a resource set greater than the size of the base set
                error(message('parallel:lang:spmd:InvalidResourceSelection', sz));
            end
        end
    end
end

% Used by ProcessRemoteResourceSet::invalidate on the workers to free a given communicator.
function iFreeMyComm()
% Use a special back-door to free communicators - the block cleanup needs to
% use them for error checking.
spmdlang.remoteBlockExecution('doFreeComm');
end

% Used by ProcessRemoteResourceSet::buildRemoteResourceSet to split the world communicator.
function val = iBuildCommHandle(newSz)
if newSz > spmdSize
    % Should never get here - two layers of previous checks would have to fail
    % to catch this...
    error(message('parallel:lang:spmd:InvalidComms', newSz, spmdSize));
end

% All workers <= spmdIndex get color 1, others get color -1 (so they get COMM_NULL)
color = (2 * double(spmdIndex <= newSz)) - 1;
key    = spmdIndex;

val    = mpiCommManip('split', color, key);
end

