# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Caching utilities."""

import inspect
import weakref


# TODO(mdan): Add a garbage collection hook for cleaning up modules.
class _TransformedFnCache(object):
  """Generic hierarchical cache for transformed functions.

  The keys are soft references (i.e. they are discarded when the key is
  destroyed) created from the source function by `_get_key`. The subkeys are
  strong references and can be any value. Typically they identify different
  kinds of transformation.
  """

  __slots__ = ('_cache',)

  def __init__(self):
    self._cache = weakref.WeakKeyDictionary()

  def _get_key(self, entity):
    raise NotImplementedError('subclasses must override')

  def has(self, entity, subkey):
    key = self._get_key(entity)
    parent = self._cache.get(key, None)
    if parent is None:
      return False
    return subkey in parent

  def __getitem__(self, entity):
    key = self._get_key(entity)
    parent = self._cache.get(key, None)
    if parent is None:
      # The bucket is initialized to support this usage:
      #   cache[key][subkey] = value
      self._cache[key] = parent = {}
    return parent

  def __len__(self):
    return len(self._cache)


class CodeObjectCache(_TransformedFnCache):
  """A function cache based on code objects.

  Code objects are good proxies for the source code of a function.

  This cache efficiently handles functions that share code objects, such as
  functions defined in a loop, bound methods, etc.

  The cache falls back to the function object, if it doesn't have a code object.
  """

  def _get_key(self, entity):
    if hasattr(entity, '__code__'):
      return entity.__code__
    else:
      return entity


class UnboundInstanceCache(_TransformedFnCache):
  """A function cache based on unbound function objects.

  Using the function for the cache key allows efficient handling of object
  methods.

  Unlike the _CodeObjectCache, this discriminates between different functions
  even if they have the same code. This is needed for decorators that may
  masquerade as another function.
  """

  def _get_key(self, entity):
    if inspect.ismethod(entity):
      return entity.__func__
    return entity