Factoring context action

gistfile1.cs

using System.Collections.Generic;
using System.Linq;
using System.Text;
using JetBrains.ReSharper.Psi.CSharp.Tree;

namespace ActiveMesa.R2P.Maths.Factor
{
  using System;
  using JetBrains.Application.Progress;
  using JetBrains.ProjectModel;
  using JetBrains.ReSharper.Feature.Services.Bulbs;
  using JetBrains.ReSharper.Feature.Services.CSharp.Bulbs;
  using JetBrains.ReSharper.Intentions.Extensibility;
  using JetBrains.ReSharper.Psi.CSharp;
  using JetBrains.TextControl;
  using JetBrains.Util;

  [ContextAction(Name = "Factor", Description = "Factors out common terms in an expression", Group = "C#")]
  public class FactorCA : ContextActionBase
  {
    private readonly ICSharpContextActionDataProvider provider;
    private Dictionary<string, int[]> histogram;
    private KeyValuePair<string, int> bestTerm;
    private int timesToExtract;
    private IAdditiveExpression root;
    private readonly HashSet<int> affectedTerms = new HashSet<int>();
    private List<List<ICSharpExpression>> flatStructure;

    public FactorCA(ICSharpContextActionDataProvider provider)
    {
      this.provider = provider;
    }

    public override bool IsAvailable(IUserDataHolder cache)
    {
      root = provider.GetSelectedElement<IAdditiveExpression>(true, true);

      // make sure we're on the root
      if (root == null) return false;
      while (root.Parent is IAdditiveExpression)
        root = (IAdditiveExpression)root.Parent;

      // get all the parts of an addition as a list
      var partsOfAddition = Flatten<IAdditiveExpression>(root).ToList();
      flatStructure = FlattenMultiplications(partsOfAddition);

      // build a frequencyTable of the most useful terms
      int elementCount = flatStructure.Count;
      histogram = new Dictionary<string, int[]>();
      for (int i = 0; i < elementCount; ++i)
      {
        var group = flatStructure[i];
        for (int index = 0; index < group.Count; index++)
        {
          string termText = group[index].GetText();
          if (!histogram.ContainsKey(termText))
            histogram.Add(termText, new int[elementCount]);
          histogram[termText][i]++;
        }
      }

      // find out the term with largest frequency
      bestTerm = histogram.ToDictionary(k => k.Key, v => v.Value.Count(z => z > 0))
                                 .ToList().OrderByDescending(x => x.Value).First();
      if (bestTerm.Value < 2) return false;

      // find out how many times we're going to take it out
      timesToExtract = histogram[bestTerm.Key].Where(z => z > 0).Min();
      var e = histogram[bestTerm.Key];
      for (int i = 0; i < elementCount; ++i)
        if (e[i] > 0)
        {
          e[i] -= timesToExtract;
          affectedTerms.Add(i);
        }
      return true;
    }

    private IEnumerable<ICSharpExpression> Flatten<T>(IBinaryExpression expr)
      where T : class, IBinaryExpression
    {
      Func<ICSharpExpression, List<ICSharpExpression>> yieldAll = t =>
      {
        var result = new List<ICSharpExpression>();
        var op = t as T;
        if (op != null)
          foreach (var e in Flatten<T>(op))
            result.Add(e);
        else result.Add(t);
        return result;
      };
      return yieldAll(expr.LeftOperand).Concat(yieldAll(expr.RightOperand));
    }

    private List<List<ICSharpExpression>> FlattenMultiplications(List<ICSharpExpression> parts)
    {
      var result = new List<List<ICSharpExpression>>();
      foreach (var part in parts)
      {
        var ops = new List<ICSharpExpression>();

        var expr = part as IMultiplicativeExpression;
        if (expr != null)
          ops.AddRange(Flatten<IMultiplicativeExpression>(expr));
        else
          ops.Add(part);

        result.Add(ops);
      }
      return result;
    }

    protected override Action<ITextControl> ExecutePsiTransaction(ISolution solution, IProgressIndicator progress)
    {
      CSharpElementFactory factory = CSharpElementFactory.GetInstance(provider.PsiModule);

      var sb = new StringBuilder();

      // first, append the first term a suitable number of times
      sb.Append(Enumerable.Repeat(bestTerm.Key, timesToExtract).Join("*"));

      // brace the terms
      sb.Append("*(");

      // (very dubious predicate lambda here)
      Action<Func<HashSet<int>, int, bool>> processTerms = predicate =>
      {
        for (int groupIndex = 0; groupIndex < flatStructure.Count; ++groupIndex)
        {
          // is this term even affected?
          if (predicate(affectedTerms, groupIndex))
          {
            // go through each of the terms that's included in this group
            foreach (string term in histogram.Keys)
              if (histogram[term][groupIndex] > 0)
              {
                // suddenly this is relevant
                sb.Append(Enumerable.Repeat(term, histogram[term][groupIndex]).Join("*"));
                sb.Append("*");
              }
            sb.RemoveLast();
            sb.Append(" + ");
          }
        }
        sb.RemoveLast(3);
      };

      // if you know a better way, let me know
      processTerms((set, i) => set.Contains(i));
      sb.Append(")");

      bool haveExtras = Enumerable.Range(0, flatStructure.Count).Any(x => !affectedTerms.Contains(x));
      if (haveExtras)
      {
        sb.Append(" + ");
        processTerms((set, i) => !set.Contains(i));
      }

      var expr = factory.CreateExpressionAsIs(sb.ToString());
      root.ReplaceBy(expr);

      return null;
    }

    public override string Text
    {
      get { return "Factor out " + Enumerable.Repeat(bestTerm.Key, timesToExtract).Join("*"); }
    }
  }
}