%MergedKeyValuesSource
% An implementation of AbstractKeyValuesSource that combines several
% instances of AbstractKeyValuesSource into a single source. The
% AbstractKeyValuesSource interface has the contract that all keys are read
% in order. This is assumed to be true in the implementation of this class.
%
% This class reads values using the steps:
%   1. It sorts the sources in ascending order of current key.
%   2. It iterates through the first N sources of the same current key.
%   3. When all of these have a new current key, rinse and repeat.

%   Copyright 2014-2020 The MathWorks, Inc.

classdef (Sealed, Hidden) MergedKeyValuesSource < parallel.internal.mapreduce.AbstractKeyValuesSource
    properties (GetAccess = private, SetAccess = immutable)
        SortFunctor; % A function handle that does the correct sort on the given type.
    end
    
    properties (Access = private)
        Keys; % A cell array of current keys.
        Sources; % A cell array of sources corresponding with keys.
        
        CurrentKeyIndex = 0; % The index into Sources of the current source being read.
        HasDataForCurrentKey = false; % A flag that indicates whether this has more data for the current key.
    end
    
    methods
        % Create a MergedKeyValuesSource from a collection of
        % instances of AbstractKeyValueSource. This object will assume
        % full control over the collection of AbstractKeyValueSource
        % instances.
        %
        % Syntax:
        %  mergedSource = MergedKeyValuesSource(source1, source2, source3);
        %
        function this = MergedKeyValuesSource(varargin)
            assert(all(cellfun(@(x)isa(x, 'parallel.internal.mapreduce.AbstractKeyValuesSource'), varargin)));
            sources = varargin(:);
            
            validSources = cellfun(@hasData, sources);
            sources(~validSources) = [];
            
            this.Keys = iParseCurrentKeys(cellfun(@nextKey, sources, 'UniformOutput', false));
            this.Sources = sources;
            this.SortFunctor = iGetSortFunctor(class(this.Keys));
        end
        
        % Returns true if and only if there is more data overall to read.
        function result = hasData(this)
            result = ~isempty(this.Sources) && any(cellfun(@hasData, this.Sources));
        end
        
        % Returns true if and only if there is more data associated with the
        % current key.
        function result = hasDataForCurrentKey(this)
            result = this.HasDataForCurrentKey;
        end
        
        % Moves the source forward to the data associated with the next
        % key. If there are no more keys, this will move the source to the
        % end.
        function key = nextKey(this)
            assert(this.hasData() && ~this.HasDataForCurrentKey);
            
            validSourceIndices = cellfun(@hasData, this.Sources(1 : this.CurrentKeyIndex));
            
            newKeys = cellfun(@nextKey, this.Sources(validSourceIndices), 'UniformOutput', ~iscell(this.Keys));
            
            this.Keys = [newKeys; this.Keys(this.CurrentKeyIndex + 1 : end)];
            this.Sources = [this.Sources(validSourceIndices); this.Sources(this.CurrentKeyIndex + 1 : end)];
            
            [this.Keys, idx] = this.SortFunctor(this.Keys);
            this.Sources = this.Sources(idx);
            
            key = this.Keys(1);
            if iscell(key)
                key = key{1};
            end
            this.CurrentKeyIndex = 1;
            this.HasDataForCurrentKey = true;
        end
        
        % Discard any remaining data associated with the current key.
        function discardRemainingDataForCurrentKey(this)
            while this.HasDataForCurrentKey
                source = this.Sources{this.CurrentKeyIndex};
                source.discardRemainingDataForCurrentKey();
                this.pNextSource();
            end
        end
        
        % Return a chunk of values up to maxLength number of entries
        % associated with the current key.
        function values = read(this, maxLength)
            assert(this.HasDataForCurrentKey);
            
            source = this.Sources{this.CurrentKeyIndex};
            values = read(source, maxLength);
            
            if ~source.hasDataForCurrentKey()
                this.pNextSource();
            end
        end
    end
    
    methods (Access = private)
        % Move to the next source object of the current key.
        function pNextSource(this)
            this.HasDataForCurrentKey = (this.CurrentKeyIndex ~= numel(this.Keys)) ...
                && isequal(this.Keys(this.CurrentKeyIndex), this.Keys(this.CurrentKeyIndex + 1));
            
            if this.HasDataForCurrentKey
                this.CurrentKeyIndex = this.CurrentKeyIndex + 1;
            end
        end
    end
end

% Parse current keys to determine whether the types are consistent and valid.
% These errors are a fail safe - they should not be hit in normal execution.
% Invalid/inconsistent keys will throw earlier in the MapReduce algorithm.
function keys = iParseCurrentKeys(inKeys)

if iscellstr(inKeys) || isstring(inKeys)
    keys = inKeys(:);
    return;
end

try
    keys = vertcat(inKeys{:});
catch
    error(message('parallel:mapreduce:InconsistentKeyType'));
end

if ~isnumeric(keys) || numel(keys) ~= numel(inKeys)
    error(message('parallel:mapreduce:InvalidKeyType', class(keys)));
end

end

function sortFunctor = iGetSortFunctor(keyType)
switch keyType
    case 'uint8'
        sortFunctor = @(keys)iCastSort(keys, 'uint8', 'int8');
    case 'uint16'
        sortFunctor = @(keys)iCastSort(keys, 'uint16', 'int16');
    case 'uint32'
        sortFunctor = @(keys)iCastSort(keys, 'uint32', 'int32');
    case 'uint64'
        sortFunctor = @(keys)iCastSort(keys, 'uint64', 'int64');
    otherwise
        sortFunctor = @sort;
end
end

function [keys, idx] = iCastSort(keys, trueType, castType)
keys = typecast(keys, castType);
[keys, idx] = sort(keys);
keys = typecast(keys, trueType);
end
