classdef StoreDB < handle_light
% The StoreDB class is a handle class to manage caching in Manopt.
%
% To create an object, call: storedb = StoreDB();
% Alternatively, call: storedb = StoreDB(storedepth); to instruct
% the database to keep at most storedepth stores in its history.
% (Note that clean up only happens when purge() is called).
%
% The storedb object is passed by reference: when it is passed to a
% function as an input, and that function modifies it, the original
% object is modified.

% This file is part of Manopt: www.manopt.org.
% Original author: Nicolas Boumal, April 3, 2015.
% Contributors: 
% Change log: 
%
%   Aug. 2, 2018 (NB)
%       Purge now selects based on time since last access, whether read or
%       write. It used to select based on last write. This should allow to
%       set a smaller value of storedepth for trustregions (in particular).
%       Also added StoreDB.remove(key) which allows to specifically remove
%       the store associated to a key. This can be used by a solver when it
%       is reasonably certain that the memory associated to a certain key
%       is no longer needed (for example, when a step was rejected in
%       trustregions or in a line-search algorithm.) Proper usage of this
%       feature should allow to set a small value for storedepth.
%       Furthermore, getNewKey() now places an empty store at the created
%       key. As a result, it is no longer legitimate to call storedb.get on
%       an unknown key: this now issues a warning, which should help debug.
%       The function removefirstifdifferent(key1, key2) only removes key1
%       if it is different from key2, which is helpful when an iteration
%       may have failed to move to a distinct point.

% TODO : protect get/setWithShared calls: limit to one, and forbid access
%        to shared memory while it has not been returned.
%        Do think of the applyStatsFun case : calls a getWithShared, does
%        not need a setWithShared. I think for statsfun there should be a
%        method "forfeitWithShared".
    
    properties(Access = public)
       
        % This memory is meant to be shared at all times. Users can modify
        % this at will. It is the same for all points x.
        shared = struct();
        
        % This memory is used by the toolbox for, e.g., automatic caching
        % and book keeping. Users should not overwrite this. It is the
        % same for all points x.
        internal = struct();
        
        % When calling purge(), only a certain number of stores will be
        % kept in 'history'. This parameter fixes that number. The most
        % recently modified stores are kept. Set to inf to keep all stores.
        storedepth = inf;
        
    end
    
    properties(Access = private)
        
        % This structure holds separate memories for individual points.
        % Use get and set to interact with this. The field name 'shared' is
        % reserved, for use with get/setWithShared.
        history = struct();
        
        % This internal counter is used to obtain unique key's for points.
        counter = uint32(0);
        
        % This internal counter is used to time calls to 'set', and hence
        % keep track of which stores in 'history' were last updated.
        timer = uint32(0);
        
    end
    
    
    methods(Access = public)
        
        % Constructor
        function storedb = StoreDB(storedepth)
            if nargin >= 1
                storedb.storedepth = storedepth;
            end
        end
        
        % Returns the store associated to a given key.
        % If the key is unknown, issues a warning and returns empty store.
        function store = get(storedb, key)
            if isfield(storedb.history, key)
                % Update access timer.
                storedb.history.(key).lastaccess__ = storedb.timer;
                % Extract the store.
                store = storedb.history.(key);
            else
                store = struct();
                store.lastaccess__ = storedb.timer;
                msg = 'Called storedb.get for a store with unknown key.';
                % If the queried key is less than the counter, it must have
                % been a valid key at some point. Inform the user that it
                % may have been purged (or removed) prematurely.
                if str2double(key(2:end)) < storedb.counter
                    msg = [msg, ' It seems that key was purged or removed.'];
                end
                msg = [msg, ' Returned an empty structure.'];
                warning('manopt:storedb:get', msg);
            end
            storedb.timer = storedb.timer + 1;
        end
        
        % Same as get, but adds the shared memory in store.shared.
        function store = getWithShared(storedb, key)
            store = storedb.get(key);
            store.shared = storedb.shared;
        end
        
        % Save the given store at the given key. If no key is provided, a
        % new key is generated for this store (i.e., it is assumed this
        % store pertains to a new point). The key is returned in all cases.
        % A field 'lastaccess__' is added/updated in the store structure,
        % keeping track of the last time that store was accessed.
        function key = set(storedb, store, key)
            if nargin < 3
                key = getNewKey(storedb);
            end
            store.lastaccess__ = storedb.timer;
            storedb.timer = storedb.timer + 1;
            storedb.history.(key) = store;
        end
        
        % Same as set, but extracts the shared memory and saves it.
        % The stored store will still have a 'shared' field, but empty.
        function key = setWithShared(storedb, store, key)
            storedb.shared = store.shared;
            store.shared = [];
            key = storedb.set(store, key);
        end
        
        % Erases a store from memory, identified by key.
        % If the key is unknown, issues a warning.
        function remove(storedb, key)
            if isfield(storedb.history, key)
                storedb.history = rmfield(storedb.history, key);
            else
                warning('manopt:storedb:remove', ...
                       ['Attempted to remove store with unknown key.\n' ...
                        'Perhaps it was purged? Try postponing purge().']);
            end
        end
        
        % Erases store at key1 if it is different from key2.
        function removefirstifdifferent(storedb, key1, key2)
            if ~strcmp(key1, key2)
                storedb.remove(key1);
            end
        end
        
        % Generates a unique key and returns it. This should be called
        % everytime a new point is generated / stored. Keys are valid field
        % names for structures. After this call, an empty store is added in
        % the cache at the newly generated key. It has minimal priority
        % vis-à-vis storedb.purge().
        function key = getNewKey(storedb)
            key = sprintf('z%d', storedb.counter);
            storedb.counter = storedb.counter + 1;
            % If we attempt to storedb.remove(key) or storedb.get(key) on
            % this newly created key before anything is stored there, we
            % will get a warning. Since there are legitimate scenarios for
            % this to happen, we place an empty store at the new key
            % immediately. Its last-access flag is set to 0 so it would be
            % the first to be purged (unless it was later accessed.)
            emptystore = struct();
            emptystore.lastaccess__ = 0;
            storedb.history.(key) = emptystore;
        end
        
        % Clear entries in storedb.history to limit memory usage.
        function purge(storedb)
            
            if isinf(storedb.storedepth)
                return;
            end
            
            if storedb.storedepth <= 0
                storedb.history = struct();
                return;
            end

            % Get list of field names (keys).
            keys = fieldnames(storedb.history);
            nkeys = length(keys);

            % If we need to remove some of the elements in the database,
            if nkeys > storedb.storedepth

                % Get the last-access counter of each element:
                % A higher number means it was accessed more recently.
                % Both read and write operations are considered 'access'.
                lastaccess = zeros(nkeys, 1, class(storedb.timer));
                for i = 1 : nkeys
                    lastaccess(i) = storedb.history.(keys{i}).lastaccess__;
                end

                % Sort the counters and determine the threshold above which
                % the field needs to be removed.
                sortlastaccess = sort(lastaccess, 1, 'descend');
                minlastaccess = sortlastaccess(storedb.storedepth);

                % Remove all fields that are too old.
                storedb.history = rmfield(storedb.history, ...
                                         keys(lastaccess < minlastaccess));
            end
            
        end % end of purge()
        
    end
    
end
