PostSharp Memoize Aspect

using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Reflection;
using PostSharp.Aspects;

/// <summary>
/// Аспект, производящий мемоизацию метода, отмечаемого данным
/// атрибутом. Методы с ref-/out-параметрами не поддерживаются.
/// </summary>
[Serializable, AttributeUsage(AttributeTargets.Method)]
public sealed class MemoizeAttribute
    : MethodInterceptionAspect, IInstanceScopedAspect
{
    /// <summary>
    /// Указывает, должен ли метод поддерживать
    /// корректную работу в многопоточной среде.
    /// </summary>
    public bool IsThreadSafe { get; set; }

    #region MemoCache

    [Serializable]
    abstract class MemoCache
    {
        public abstract bool TryResolve(
            object arg, out object result);
        public abstract void AppendItem(
            Arguments arg, int index, object result);
    }

    [Serializable]
    abstract class MemoCache<T, TResult> : MemoCache
    {
        static readonly Func<MemoCache> NestedCacheFactory;

        static MemoCache()
        {
            // если элементами кэша данного типа
            // являются другие вложенные кэши
            if (typeof(TResult).IsSubclassOf(typeof(MemoCache)))
            {
                NestedCacheFactory =
                    GetCacheFactory(typeof(TResult));
            }
        }

        protected abstract void AppendImpl(T arg, TResult result);

        public sealed override void AppendItem(
            Arguments arg, int index, object result)
        {
            if (NestedCacheFactory == null)
            {
                // тривиально добавляем в кэш
                AppendImpl((T) arg[index], (TResult) result);
            }
            else
            {
                // создаём экземпляр вложенного кэша
                var nested = NestedCacheFactory();

                AppendImpl( // и добавляем его в кэш
                    (T) arg[index],
                    (TResult) (object) nested);

                // кэшируем следующий аргумент
                nested.AppendItem(arg, index + 1, result);
            }
        }
    }

    /// <summary>
    /// Вариант кэша на базе обычного словаря.
    /// </summary>
    [Serializable]
    sealed class DictionaryCache<T, TResult>
                        : MemoCache<T, TResult>
    {
        readonly Dictionary<T, TResult> cache
            = new Dictionary<T, TResult>();

        public static MemoCache CreateInstance()
        {
            return new DictionaryCache<T, TResult>();
        }

        public override bool TryResolve(object arg, out object result)
        {
            TResult value;
            if (cache.TryGetValue((T) arg, out value))
            {
                result = value;
                return true;
            }
            else
            {
                result = null;
                return false;
            }
        }

        protected override void AppendImpl(T arg, TResult result)
        {
            cache.Add(arg, result);
        }
    }

    /// <summary>
    /// Вариант кэша на базе конкурентного словаря.
    /// </summary>
    [Serializable]
    sealed class ConcurrentCache<T, TResult>
                        : MemoCache<T, TResult>
    {
        readonly ConcurrentDictionary<T, TResult> cache
            = new ConcurrentDictionary<T, TResult>();

        public static MemoCache CreateInstance()
        {
            return new ConcurrentCache<T, TResult>();
        }

        public override bool TryResolve(object arg, out object result)
        {
            TResult value;
            if (cache.TryGetValue((T) arg, out value))
            {
                result = value;
                return true;
            }
            else
            {
                result = null;
                return false;
            }
        }

        protected override void AppendImpl(T arg, TResult result)
        {
            cache.AddOrUpdate(arg, result, (_, x) => x);
        }
    }

    /// <summary>
    /// Возвращает фабрику создания экземпляров кэша по типу.
    /// </summary>
    static Func<MemoCache> GetCacheFactory(Type cacheType)
    {
        // ищем метод "public static MemoCache CreateInstance()"
        var methodInfo = cacheType.GetMethod(
            "CreateInstance", BindingFlags.Static | BindingFlags.Public);

        // и создаём из него делегат для быстрого создания экземпляров
        return (Func<MemoCache>)
            Delegate.CreateDelegate(typeof(Func<MemoCache>), methodInfo);
    }

    /// <summary>
    /// Создаёт тип кэша, соответствующий типам параметров
    /// заданного метода и требованиям к многопоточной работе.
    /// </summary>
    Type GetRootCacheType(MethodInfo method)
    {
        Debug.Assert(method != null);

        var parameters = method.GetParameters();
        var resultType = method.ReturnType;

        // определяем тип используемого кэша
        var cacheType = IsThreadSafe
            ? typeof(ConcurrentCache<,>)
            : typeof(DictionaryCache<,>);

        // перебираем параметры с конца
        for (int i = parameters.Length - 1; i >= 0; i--)
        {
            // формируем тип "Cache<T1, Cache<T2, Cache<T3, TResult>>>",
            // в котором типы T1, T2, T3 соответствуют параметрам метода:
            resultType = cacheType.MakeGenericType(
                parameters[i].ParameterType, resultType);
        }

        return resultType;
    }

    #endregion
    #region Methods

    /// <summary>
    /// Процедура валидации использования аспекта мемоизации.
    /// </summary>
    public override bool CompileTimeValidate(MethodBase method)
    {
        var mi = (MethodInfo) method;
        if (mi.ReturnType == typeof(void))
        {
            throw new InvalidOperationException(
                "Аспект следует применять только " +
                "на методы, возвращающие значение.");
        }

        var paremeters = mi.GetParameters();
        if (paremeters.Length == 0)
        {
            throw new InvalidOperationException(
                "Аспект следует применять только " +
                "на методы, имеющие параметры.");
        }

        foreach (var parameter in paremeters)
        {
            if (parameter.IsIn || parameter.IsOut)
            {
                throw new InvalidOperationException(
                    "Аспект невозможно использовать с методами, " +
                    "обладающими ref-/out-параметрами.");
            }
        }

        return true;
    }

    MemoCache cacheRoot;

    /// <summary>
    /// Обработчик вызова мемоизируемого метода.
    /// </summary>
    public override void OnInvoke(MethodInterceptionArgs args)
    {
        MemoCache argCache = this.cacheRoot;
        Arguments arguments = args.Arguments;
        object result = null;
        int index = 0;

    LookupArg: // последовательно извлекаем значения из кэшей
        if (argCache.TryResolve(arguments[index++], out result))
        {
            // если не последний аргумент, то кэш
            if (index < arguments.Count)
            {
                argCache = (MemoCache) result;
                goto LookupArg; // да, это goto!
            }

            args.ReturnValue = result;
        }
        else // промах кэша, вызываем метод и кэшируем
        {
            args.Proceed();
            argCache.AppendItem(
                arguments, index - 1, args.ReturnValue);
        }
    }

    static Func<MemoCache> RootCacheFactory;

    /// <summary>
    /// Статическая инициализация аспекта,
    /// создаёт кэш для мемоизации статических методов.
    /// </summary>
    public override void RuntimeInitialize(MethodBase method)
    {
        var type = GetRootCacheType((MethodInfo) method);
        RootCacheFactory = GetCacheFactory(type);

        if (method.IsStatic)
            this.cacheRoot = RootCacheFactory();

        base.RuntimeInitialize(method);
    }

    /// <summary>
    /// Создание экземпляра аспкета уровня экземпляра,
    /// попадает в конструктор типа, экземплярный метод
    /// которого подвергается аспекту мемоизации.
    /// </summary>
    public object CreateInstance(AdviceArgs adviceArgs)
    {
        return new MemoizeAttribute {
            IsThreadSafe = this.IsThreadSafe
        };
    }

    /// <summary>
    /// Инициализация аспекта уровня экземпляра.
    /// </summary>
    public void RuntimeInitializeInstance()
    {
        this.cacheRoot = RootCacheFactory();
    }

    #endregion
}

class Foo
{
    [Memoize]
    static int StaticFact(int x)
    {
        Console.WriteLine("=> StaticFact({0}) call", x);

        if (x == 0) return 1;
        return x * StaticFact(x - 1);
    }

    [Memoize(IsThreadSafe=true)]
    int InstanceFact(int x)
    {
        Console.WriteLine("=> InstanceFact({0}) call", x);

        return Enumerable
            .Range(1, x)
            .Aggregate(1, (a, b) => a * b);
    }

    static void Main()
    {
        Console.WriteLine("SataticFact(2) = {0}", StaticFact(2));
        Console.WriteLine("SataticFact(2) = {0}", StaticFact(2));
        Console.WriteLine("SataticFact(7) = {0}", StaticFact(7));
        Console.WriteLine("SataticFact(7) = {0}", StaticFact(7));

        Console.WriteLine();

        var a = new Foo();
        var b = new Foo();

        Console.WriteLine("a.InstanceFact(4) = {0}", a.InstanceFact(7));
        Console.WriteLine("a.InstanceFact(4) = {0}", a.InstanceFact(7));

        Console.WriteLine("b.InstanceFact(2) = {0}", b.InstanceFact(7));
        Console.WriteLine("b.InstanceFact(2) = {0}", b.InstanceFact(7));
    }
}