anonymous / gist:4326509

Created December 18, 2012

using System.Collections.Generic;
using System.Diagnostics;
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
  {
    enum Sign
    {
      Plus,
      Minus
    }
 
    [DebuggerDisplay("Expr={Expression.GetText()},Sign={Sign}")]
    private class CSharpExpressionWithSign
    {
      public CSharpExpressionWithSign()
      {
        Sign = Sign.Plus;
      }
 
      public ICSharpExpression Expression { get; set; }
      public Sign Sign { get; set; }
    }
 
    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<CSharpExpressionWithSign>> 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 = FlattenAdditions(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].Expression.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<CSharpExpressionWithSign> FlattenAdditions(IAdditiveExpression expr)
    {
      Func<ICSharpExpression, List<CSharpExpressionWithSign>> yieldAll = t =>
        {
          var result = new List<CSharpExpressionWithSign>();
          var op = t as IAdditiveExpression;
          if (op != null)
            foreach (var e in FlattenAdditions(op))
              result.Add(e);
          else result.Add(new CSharpExpressionWithSign { Expression = t });
          return result;
        };
 
      var left = yieldAll(expr.LeftOperand);
      var right = yieldAll(expr.RightOperand);
      if (right.Any() && expr.IsSubtraction())
        right[0].Sign = Sign.Minus;
      return left.Concat(right);
    }
 
    private IEnumerable<CSharpExpressionWithSign> Flatten<T>(IBinaryExpression expr)
      where T : class, IBinaryExpression
    {
      Func<ICSharpExpression, List<CSharpExpressionWithSign>> yieldAll = t =>
      {
        var result = new List<CSharpExpressionWithSign>();
        var op = t as T;
        if (op != null)
          foreach (var e in Flatten<T>(op))
            result.Add(e);
        else result.Add(new CSharpExpressionWithSign { Expression = t });
        return result;
      };
      return yieldAll(expr.LeftOperand).Concat(yieldAll(expr.RightOperand));
    }
 
    private List<List<CSharpExpressionWithSign>> FlattenMultiplications(List<CSharpExpressionWithSign> parts)
    {
      var result = new List<List<CSharpExpressionWithSign>>();
      foreach (var part in parts)
      {
        var ops = new List<CSharpExpressionWithSign>();
 
        var expr = part.Expression as IMultiplicativeExpression;
        if (expr != null && expr.IsMultiplication())
        {
          var flattened = Flatten<IMultiplicativeExpression>(expr).ToList();
          if (flattened.Any())
            flattened[0].Sign = part.Sign;
          ops.AddRange(flattened);
        }
      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 =>
        {
          bool first = true;
          var terms = Enumerable.Range(0, flatStructure.Count).Where(i => predicate(affectedTerms, i)).ToList();
          for (int i = 0; i < terms.Count; i++)
          {
            var groupIndex = terms[i];
            
            // check if a negative sign is necessary
            bool isNegative = flatStructure[i].Any(x => x.Sign == Sign.Minus);
            if (isNegative)
              sb.Append("-");
            if (i > 0 && !isNegative)
              sb.Append(" + ");
 
            // go through each of the terms that's included in this group
            bool any = false;
            foreach (string term in histogram.Keys.OrderBy(k => k))
              if (histogram[term][groupIndex] > 0)
              {
                // suddenly this is relevant
                sb.Append(Enumerable.Repeat(term, histogram[term][groupIndex]).Join("*"));
                sb.Append("*");
                any = true;
              }
            if (any) sb.RemoveLast();
            else sb.Append("1");
          }
        };
 
      // 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("*"); }
    }
  }
}