Metaprogramming: creating compiled functions from inter-dependent code blocks
To compose code from blocks you could use simple quoting mechanism like following:
ClearAll[quote, unquote, eval]
SetAttributes[{quote, unquote}, HoldAllComplete]
quote@expr : Except@_Symbol :=
Unevaluated@expr /. {x : _unquote | _quote :> x, s_Symbol -> quote@s}
unquote@args___ := args
eval = # /. HoldPattern@quote@s_ :> s &;
or its more elaborate version.
Those code fragments that are intended to appear in final compiled code should be quoted, and those that should evaluate, during "code assembly", should be unquoted:
ClearAll@f
f[x_, g_] := Module[{code, block},
block = quote[z = y x];
code = quote@Compile[{},
quote[Module][{y, z}, y = g[x]; unquote@block; z],
CompilationTarget -> "C"
];
code // eval
];
In above definition constituent code blocks are quoted. block
variable inside Compile
is unquoted because we want it to evaluate, so that quoted code, assigned to this variable, will be inserted in desired position. There's also quote
around Module
which at first might seem superfluous, but it's needed to prevent outer Module
(one with code
and block
variables) from renaming variables of inner Module
, so that y
and z
used in block
and those used in code
are interpreted as the same variables.
To make sure we're actually compiling what we want, we can inspect final compiled expression by evaluating f
function in environment with Compiled
dynamically replaced by Hold
:
Block[{Compile = Hold}, f[3, (#^2 &)]]
(* Hold[{}, Module[{y, z}, y = (#1^2 &)[3]; z = y 3; z], CompilationTarget -> "C"] *)
We can inspect compiled code:
Needs@"CompiledFunctionTools`"
f[3, (#^2 &)] // CompilePrint
(*
No argument
3 Integer registers
Underflow checking off
Overflow checking off
Integer overflow checking on
RuntimeAttributes -> {}
I0 = 3
Result = I2
1 I1 = Square[ I0]
2 I2 = I1 * I0
3 Return
*)
Evaluation of compiled function gives expected result:
f[3, (#^2 &)][]
(* 27 *)
You already have a very good and general answer for your question. Just for completeness I wanted to show you an alternative which is much less general but shows that there is actually not too much magic in manipulating unevaluated code:
f[x_, g_] := Module[{template, code, block, module},
template = HoldComplete[
Compile[{}, module[{y, z}, y = g[x]; block; z], CompilationTarget -> "C"]
];
code = template /. block :> (z = y*x) /. module -> Module;
ReleaseHold[code]
];
The idea is straightforward: define a template expression for your code, then use standard pattern matching to insert code fragments into that code (and implicitly also use the standard litteral insertion of function arguments) and when done evaluate the result using ReleaseHold
. The only part that is a bit tricky is that you need to prevent the outer Module
to rename the inner Module
s local variables (exactly the same reason why jkuczm needed an extra quote for Module
). Here I solved it by using a dummy symbol module
which is only replaced by Module
after the outer Module
has done its magic.
One technique that you might need when going that way is how you can construct a replacement rule from a held expression. One possibility is this:
block = Hold[z = y*x];
blockrule = RuleDelayed @@ Prepend[block, HoldPattern[block]]
template /. blockrule /. module :> Module
I came up with my own way of dealing with this. I'm not sure if it's the simplest way. I was careful to make sure I dealt with all possible namespace issues, as well as possible recursive calls where code blocks need to be inserted within code blocks. My method also makes sure to avoid evaluating the arguments of any functions until the full code has been assembled.
I define two functions intended for the end user to use directly, Meta[]
and Macro[]
. In addition to these are a few helper functions, and some specialized versions of Hold[]
(used so that they can be found in patterns, without conflicting with any literal Hold[]
s the user may have typed in the code in the first place!).
(* Macro is designed to access the source code of a Function body and return it
wrapped in Hold, with the appropriate variables plugged in *)
SetAttributes[Macro,HoldAll];
(* One trivial definition so that the blue highlighting goes away *)
Identity[Macro[any___]]^:=Macro[any];
(* We create a specialized version of Hold for dealing with Macros, so that we don't
erroneously release other Holds that the programmer might have
actually intended to stay! *)
SetAttributes[HoldMacro,HoldAll];
Identity[HoldMacro[any___]]^:=HoldMacro[any];
(* Various specialized Hold functions for destructuring a function call and
reassembling it, without prematurely evaluating its arguments
(just in case an argument is a global variable that has been localized!) *)
SetAttributes[HoldArg,HoldAll];
Identity[HoldArg[any___]]^:=HoldArg[any];
(* We leave Macro inert and instead define a function that accesses the contents.
Notice that the Macro head already has the HoldAll attribute;
thus getMacroSource does not need HoldAll. In fact, we *want* it to evaluate! *)
getMacroSource[Macro[fun_[args___]]]:=
With[{def=OwnValues[fun][[1,2]],heldargs=ReleaseHold@(HoldArg/@Hold[args])},
If[Head[def]===Function,
Replace[MapAt[HoldMacro,def,{2}][heldargs],HoldArg[arg_]:>arg,Infinity]
]];
(* Default case *)
getMacroSource[expr:Except[_Macro]]:=expr;
(* Replace any macros in a block of code with their source code, wrapped in HoldMacro *)
SetAttributes[injectMacros,HoldAll];
injectMacros[expr_]:=
(* search for macros in reverse order of depth *)
With[{positions=Reverse@SortBy[Position[expr,HoldPattern@Macro[_]],Length]},
If[Length[positions]>0,
(* Traverse the list of macro positions *)
Fold[
With[{source=getMacroSource@Extract[#1,#2]},
(* Recurse through the macro source to look for more macros *)
(* Always inserts the deepest nested macros first! *)
ReplacePart[#1,injectMacros[source],#2]]&,
expr,
positions],
(* else *)
(* Make sure there is an exit from recursion *)
expr]
]
(* Meta can be used with any expression, but especially as Meta@Compile[...].
Meta keeps its argument wrapped in Hold, looks for any unexpanded Macro heads,
and expands them to their source code before evaluating the entire expression *)
SetAttributes[Meta,HoldFirst];
(* Expand Macro heads to source code wrapped in Hold with appropriate values plugged in;
then remove the inner Hold heads to inject the source code.
Finally remove the outer Hold *)
Meta[expr_]:=
ReleaseHold@
Replace[
injectMacros[Hold[expr]],
{HoldMacro[any___]:>any},
Infinity]
The usage is as follows:
Needs["CompiledFunctionTools`"];
bar = Function[{x, y}, x + y];
foo = Meta@Compile[{{x, _Integer}, {y, _Integer}},
x*y*Macro@bar[x, x] + Macro@bar[y, x]];
Notice that bar
is defined as an ordinary function; its usage as a macro in the source code for foo
happens at the time you write the code for foo
. We can see the source code that Compile[]
actually sees:
Extract[foo, {{-2}}]
(* {Function[{x, y}, x y (x + x) + (y + x)]} *)
which clearly contains no function calls to bar[]
, but has completely inlined the code. bar[]
can also contain scoping constructs like Module
and they will behave correctly. bar[]
can also contain additional macros, and their code will also be expanded.
I'm not sure if this is the simplest way to achieve the desired result, as I see one of the solutions above just uses Rules directly. But I think I remember running into subtle problems with scoping constructs and function arguments prematurely evaluating. The objective here is to make sure everything is held until it has been put in place, so that Compile[]
sees it literally as though it had been typed all at once.