(* :Name: OptimizeExpression` *)

(* :Title: Expression Optimization. *)

(* :Author: Mark Sofroniou *)

(* :Summary:
 This package adds the procedure OptimizeExpression for performing common
 sub-expression optimization on Mathematica expressions.
 An new data structure OptimizedExpression is introduced.
 Several elements of the system, such as the compiler, know how to
 deal with OptimizedExpressions.
 The function CostExpression measures the computational expense of numerical
 evaluation of an expression.
 The additional package Format.m (MathSource) enables the production of
 efficient C and Fortran code. *)

(* :Context: NumericalMath` *)

(* :Package Version: 1.0 *)

(* :Copyright: Copyright 1997-1999, Wolfram Research, Inc. *)

(* :History:
 Updated March 1997.
 Rewritten May 1996.
 (Original Version by Mark Sofroniou, September 1993.) *)

(* :Keywords:
 Compile, CostExpression, OptimizeExpression, Expression Optimization,
 Polynomial, Sub-Expression, Syntactic. *)

(* :Source:
 Book: Elements of Mathematica Programming, edited by Troels Petersen. *)

(* :Mathematica Version: 4.0 *)

(* :Limitations:
 This package enables syntactic optimization in linear time
 (O(n) operations).
 The term optimization is used in the sense of improving the arithmetic
 operation count, rather than the best possible.
 Expressions containing non-binary arithmetic operators (Plus
 and Times) are optimized by matching only entire sub-expressions
 without extracting numeric coefficients. These operations rarely
 dominate the computational cost.
 Options enable control over the optimization process.
 The numerical stability of expressions is largely unaffected
 by the optimization process. *)

BeginPackage["NumericalMath`OptimizeExpression`"]

CostExpression::usage = "CostExpression[expr, opts] returns the sum of
operations present in expr, including basic arithmetic operations and
function calls. Options opts control the type of operation count
performed."

OptimizeExpression::usage = "OptimizeExpression[expr, opts] transforms expr
into a list of optimization statements and an optimized expression wrapped in
OptimizedExpression. Several elements of the system know how to deal
with OptimizedExpression objects including the compiler.
Optimization is performed in linear time (O(n) operations)
providing an efficient means of reducing the arithmetic operation count - not
the best possible. The optimization performed is mainly syntactic (only
exact common sub-expressions are matched) with a few heuristics. Options
opts control the type of optimizations performed and the format of the
output."

OptimizedExpression::usage = "OptimizedExpression is a data container
introduced by OptimizeExpression for representing the optimized
directed acyclic graph of an expression. Function, Compile and SetDelayed
all know about OptimizedExpression and convert it to an appropriate form for
subsequent computations. Normal can be used to recover the original
expression."

(* Options: *)

Binary::usage = "Binary is a setting of the option OptimizePowers,
specifying whether to perform repeated binary factoring of exponents."

OptimizeLevel::usage = "OptimizeLevel is an option of OptimizeExpression
specifying the depth to which optimizations are performed in an expression."

OptimizePowers::usage = "OptimizePowers is an option of OptimizeExpression
specifying whether to optimize sub-expressions with head Power.
Integer and rational powers can be factored efficiently by repeated
binary exponentiation. OptimizePowers may evaluate to Binary, Literal
or False."

OptimizeVariable::usage = "OptimizeVariable is an option of OptimizeExpression
specifying a pair {symb,form} where symb is the optimization variable
to introduce and form is the type of format to use.
The format can be either Sequence (o1,o2,...) or Array (o[1],o[2],...).
Consecutively numbered variables are returned."

CostAsBinary::usage = "CostAsBinary is an option of CostExpression specifying
a list of symbols which should be treated as binary operators.
CostAsBinary may evaluate to a list of Symbols or None."

CostLevel::usage = "CostLevel is an option of CostExpression
specifying the depth to which an operation count is performed
in an expression."

CostPowers::usage = "CostPowers is an option of CostExpression specifying
whether to consider x^y as being calculated as Exp[y*Log[x]]. The setting
Integer assumes in addition that integer powers are calculated using
(y-1) multiplications, which is useful when costing expressions optimized
using the option OptimizePowers->Binary. CostPower may evaluate to
Integer, True or False. See also CostDivisions."

(* Default optimization symbol. *)

$::usage="$ is the default optimization symbol introduced by
OptimizeExpression."

Unprotect[Binary, CostExpression, CostDivisions, CostExp, CostLevel,
  CostPowers, ExcludedForms, OptimizeExpression, OptimizedExpression,
  OptimizePowers, OptimizeVariable, $ ];

Begin["`Private`"]

(* Function to check the data types of CostExpression options with error
messages. *)

Options[CostExpression] = {
  CostAsBinary -> {Plus, Times}, CostLevel -> -2, CostPowers -> False,
  ExcludedForms -> {_?AtomQ, _CompoundExpression, _List,
    _OptimizedExpression}
  };

CostExpression::opts = "The option `1` did not evaluate to `2`.";

pcerrmsgs = {
  {"option CostAsBinary", "a list of symbols or None"},
  {"option CostLevel", "a negative integer or -Infinity"},
  {"option CostPowers", "Integer, True or False"},
  {"option ExcludedForms", "a list of expressions"} };

CostOptTest[opts___]:=
  Module[{defaults, optvals, types},

    defaults = Options[CostExpression];

    optvals = Map[First, defaults] /. {opts} /. defaults;

    types = {
      MatchQ[optvals[[1]], None|{___Symbol}], (* Binary operators *)
	  MatchQ[optvals[[2]], _Integer?Negative|-Infinity], (* Level *)
      MatchQ[optvals[[3]], Integer|True|False], (* Power *)
      MatchQ[optvals[[4]], _List] (* Excluded forms *)
      };

    Check[
      MapThread[
        If[ #1, #1, Message[CostExpression::opts, Apply[Sequence,#2]] ]&,
        {types, pcerrmsgs}
      ]; optvals,          (* Return list of option values. *)
      $Failed,             (* Option of wrong type. *)
      CostExpression::opts (* Check only for these messages. *)
    ]
  ];

CostExpression[expr_,opts___?OptionQ]:=
  Module[{optvals},
    MainCost[expr,optvals] /;
      (optvals = CostOptTest[opts])=!=$Failed
  ];

(* Function for counting cost of basic operations. *)

MainCost[expr_,{costbins_, costlevel_, costpower_, exclude_}]:=
  Module[{CostFunction, fvar},

(* Prevent evaluation during costing. *)

    SetAttributes[{CostFunction}, HoldAll];

(* Don't count specified patterns. *)

    Map[ Function[ fvar, CostFunction[fvar]:= 0 ], exclude ];

(* Binary operators. *)

    If[ costbins=!=None,
      Map[
        Function[
          fvar,
          CostFunction[Pattern[x, Blank[fvar]]]:= (Length[Unevaluated[x]]-1) fvar
        ],
        costbins
      ]
    ];

(* Rules to cost powers. *)

    HoldPattern[CostFunction[E^x_]]:= Exp + CostExpression[Unevaluated[x]];

(* Integer powers as multiplications and divisions. *)

    If[SameQ[costpower, Integer],
      HoldPattern[CostFunction[_^y_Integer]]:= (Abs[y]-1) Times
    ];

(* x^y calculated as Exp[y*Log[x]]. *)

    If[MatchQ[costpower, Integer|True],
      CostFunction[_Power]:= Exp+Log+Times;
    ];

(* Cost remaining functions. *)

    CostFunction[x_]:= Head[Unevaluated[x]];

(* Cost required operators. *)

    Apply[
      Plus,
      Map[CostFunction, Level[Unevaluated[{expr}], costlevel, Unevaluated] ]
    ]

  ]; (* End of MainCost *)


(* Optimization of expressions. *)

(* Function to check the data types of OptimizeExpression options with
 error messages. *)

Options[OptimizeExpression] = {ExcludedForms -> {_?AtomQ, _List},
  OptimizeLevel -> -2, OptimizePowers -> Binary,
  OptimizeVariable -> {$, Sequence}};

OptimizeExpression::opts = "The option `1` did not evaluate to `2`.";

opterrmsgs = {
    {"option ExcludedForms", "a list of expressions"},
    {"option OptimizeLevel", "a negative integer or -Infinity"},
    {"option OptimizePowers", "Binary, Literal or False"},
    {"option OptimizeVariable", "a pair of the form {Symbol, Sequence|Array}"}
  };

OptimizeOptionTest[opts___]:= 
  Module[{defaults, optvals, types},

    defaults = Options[OptimizeExpression];

    optvals = Map[First, defaults] /. {opts} /. defaults;

    types = {
      MatchQ[optvals[[1]], _List], (* Excluded forms *)
      MatchQ[optvals[[2]], _Integer?Negative|-Infinity], (* Level *)
      MatchQ[optvals[[3]], Literal|Binary|False], (* Powers *)
      MatchQ[optvals[[4]], {_Symbol,Sequence|Array}]}; (* Variable, type *)

    Check[
      MapThread[
        If[ #1, #1, Message[OptimizeExpression::opts, Apply[Sequence, #2]] ]&,
        {types, opterrmsgs}
      ]; optvals,      (* Return list of option values. *)
      $Failed,         (* Option of wrong type. *)
      OptimizeExpression::opts   (* Check only for these messages. *)
    ]
  ]; (* End of OptimizeOptionTest. *)


(* Invoke option testing. *)

OptimizeExpression[expr_,opts___?OptionQ]:=
  Module[{optvals},
    MainOptimize[expr,optvals] /;
      (optvals = OptimizeOptionTest[opts])=!=$Failed
  ];

(* Define optimization process. *)

MainOptimize[ expr_, {exclude_, level_, optpower_, {outvar_, outform_}} ]:=
  Module[{BinDecomposition, downvals, fvar, index = 0, MakeBinRule, optexpr,
      OptRule, optvar},
    Block[{BinDecomposition, downvals, MakeBinRule, optexpr, OptRule, optvar},

    OptRule[e_optvar]:= e;

(* Don't make rules for specified patterns. Include the optimization variable
 itself in the list of excluded forms. *)

    Map[ Function[ fvar, OptRule[e:fvar]:= e], exclude ];

(* Rules for expressions with head Power. *)

    If[ MatchQ[optpower,Binary|Literal],
      HoldPattern[OptRule[e:(_^-1)]]:= (* Rules for reciprocals. *)
        Function[ fvar, OptRule[e]:= OptRule[e] = fvar; fvar ] @
          optvar[++index];
      HoldPattern[OptRule[x_^y_?Negative]]:= OptRule[ OptRule[x^(-y)]^-1 ];
    ];

(* Rules to binary factor positive integer and rational powers. *)

    If[SameQ[optpower,Binary], 

(* Save all binary power optimizations (even if they only occur once). *)

      MakeBinRule[e_]:= OptRule[e]=optvar[++index];

(* Decompose as products of binary powers. *)

      HoldPattern[OptRule[x_^p_Integer]]:=
        OptRule[ Apply[ Times, BinDecomposition[ x, BinExponents[p] ] ] ];

      SetAttributes[BinDecomposition,{Listable}];
      BinDecomposition[x_,1]:= x;    (* Stopping criterion. *)

(* Store binary power decomposition and recurse down one level. *)

      BinDecomposition[x_,p_]:= 
        BinDecomposition[ x, p ] = MakeBinRule[ BinDecomposition[ x, p-1]^2 ];

(* Rule for positive rational power with unit numerator. *)

      HoldPattern[OptRule[x:(_^Rational[1, _])]]:= MakeBinRule[x];

(* Rule to binary decompose rational power as: x^(m/n) -> (x^(1/n))^m. *)

      HoldPattern[OptRule[x_^Rational[m_, n_]]] /; m<n :=
          OptRule[ OptRule[x^Rational[1, n]]^m ];

(* Rule to decompose rational power as a product
  x^(m/n) -> x^(q+r/n) -> x^(q)*x^(r/n) -> x^(q)*(x^(1/n))^r. *)

      HoldPattern[OptRule[x_^Rational[m_, n_]]]:=
        OptRule[
          OptRule[ x^Quotient[m, n]]*
          OptRule[ OptRule[x^Rational[1, n]]^Mod[m, n] ]
        ];

    ];

(* General rule for functions, non-numeric and non-binary powers etc.
 * Stores a sub-expression by making a new DownValue for OptRule and replacing
 * the sub-expression by a new optimization variable. *)

    OptRule[e_]:=
      Function[ fvar, OptRule[e]:= OptRule[e] = fvar; fvar ] @
        optvar[++index];

(* Map function for creating optimization rules through the specified levels.
 The outer application is required for binary power factorization. *)

    optexpr = OptRule[ Map[ OptRule, expr, level ] ];

(* Find occurrences of temporary variables in optimized expression and
 in optimization rules. *)

    downvals = DownValues[OptRule];

(* Decide which temporary variables to keep and which to discard by pattern
 look-up. Discard rules by assigning optimization variable to subexpression. *)

    Cases[
      downvals,
      HoldPattern[_[rhs_]=lhs_]:>(lhs=rhs),
      {2}
    ];

(* Save ordered list of repeated rules. *)

    downvals =
      Sort[
        Cases[
          downvals,
          HoldPattern[_[_[rhs_]]:>lhs_optvar]:>{lhs,rhs}
        ]
      ];

(* Exit if no optimization performed. *)

    If[downvals==={}, Return[expr] ];

(* Format optimization variable. Reset index for consecutively numbered
 output variables. *)

    index = 0;

    If[outform===Sequence,

(* Write optimization variable as consecutive symbols $1, $2, ...
 Store indexing string function (not localized). *)

      Function[
        fvar,
        optvar[i_]:= optvar[i]= ToExpression[fvar<>IndexString[++index]]
      ] @ ToString[outvar],

(* Write optimization variable as a consecutive array $[1], $[2], ... *)

      optvar[i_]:= optvar[i]= outvar[++index]
    ];

(* Create an optimized object, apply formatting rules for optimization
 variable and substitute discarded rules into optimized expression. *)

    MakeProc[ outvar, outform, Thread[downvals], optexpr ]

    ]
  ]; (* End of MainOptimize.*)


(* Calculate a product of binary power decompositions of an expression.
 Binary power decomposition can be constructed as 2^(BinExponents[p]-1) *)

BinExponents[p_]:= BinExponents[p] =
  Flatten[Position[Reverse[IntegerDigits[p,2]],1]];

(* This definition is used to convert an index to a string for use by
 * sequential variables $1, $2 etc. The definition is not localised so
 * that the results can be reused in subsequent computations. *)

IndexString[i_]:= IndexString[i]=ToString[i];


(* Convert output to a procedure. Make sure OptimizedExpression is not the head
 of this rule to avoid conflict with the later rule for SetDelayed. *)

MakeProc[ _, Sequence, {locals_, optseq_}, optexpr_]:=
  Block[{Set},
  	BuildProc[locals, Evaluate[Thread[ Set[locals,optseq] ]], optexpr]
 ];

MakeProc[ localsymbol_, Array, {locals_, optseq_}, optexpr_]:=
  Block[{Set},
      BuildProc[{localsymbol}, Thread[ Set[locals,optseq] ], optexpr]
	];

BuildProc[locals_, {body__}, optexpr_]:=
  Block[{Set},
    OptimizedExpression[ Module[ locals, body; optexpr] ]
  ];

(* Definitions for system integration. *)

(* SetDelayed *)

OptimizedExpression /:
  HoldPattern[ SetDelayed[ lhs_, OptimizedExpression[expr_] ] ]:=
    lhs:=expr;

(* Compile *)

OptimizedExpression /:
  Compile[args_, OptimizedExpression[expr_], info___]:=
    Compile[ args, expr, info ];

(* Functions and pure functions *)

OptimizedExpression /:
  Function[ OptimizedExpression[expr_], attr___]:=
    Function[ expr, attr ];

OptimizedExpression /:
  Function[args_, OptimizedExpression[expr_], attr___]:=
    Function[ args, expr, attr ];

(* Numerical approximation, NValues etc. *)

N[ OptimizedExpression[expr_] ]:= N[ Unevaluated[expr] ];

N[ OptimizedExpression[expr_], prec_ ]:= N[ Unevaluated[expr], prec ];

OptimizedExpression /:
  h:(SetAccuracy|SetPrecision)[ OptimizedExpression[expr_],
       prec:(_?Positive|Infinity)]:=
    h[ Unevaluated[expr], prec ];

(* Normal - no definition for OptimizeExpression since Normal does not have a
  Hold attribute. *)

OptimizedExpression /:
  Normal[OptimizedExpression[expr_]]:= expr;

(* Iterators. Turn the optimized procedure into a lexically
 * scoped function of the iterator variables. *)

OptimizedExpression /:
  (h:Alternatives[Do,Product,Sum,Table])[
     OptimizedExpression[expr_], itargs__]:=
       With[{fargs = First[ Thread[{itargs}] ]},
         h[ Function[ fargs, expr][fargs], itargs ]
       ];

End[];  (* End `Private` Context. *)

(* Protect exported symbols. *)

SetAttributes[{OptimizeExpression, CostExpression}, ReadProtected];
SetAttributes[OptimizedExpression, {HoldAll, NHoldAll, ReadProtected}];

Protect[Binary, CostExpression, CostDivisions, CostExp, CostLevel,
  CostPowers, ExcludedForms, OptimizeExpression, OptimizedExpression,
  OptimizePowers, OptimizeVariable, $ ]; 

EndPackage[];    (* End package Context. *)

