Garbage collection for lexical closures

Question

Mathematica is known not to garbage collect temporary variables in certain cases. I am particularly interested in the very useful programming pattern known as a lexical closure. A somewhat trivial example is

getMultiplier[n_] := Module[{f},
  f[x_] := n*x;
  f
];

Defining M = getMultiplier[4] returns a function that multiplies by 4, for example M[5] will return 20. But if one removes the symbol M by calling Remove[M] then the temporary variable f that it referred to will remain in memory and not be garbage collected. This is also true with $HistoryLength = 0.

Question: What workarounds can be used to avoid the memory leak that comes with this programming pattern (lexical closure) in Mathematica? I include two workarounds, but I hope somebody has a more general and more robust solution.

Workaround 1: Avoid generating temporary variables in the first place, using an anonymous function as in

getMultiplier[n_] := (n*#&);

This workaround has various limitations, for example it prohibits common Mathematica patterns such as memoization f[x_] := f[x] = n*x.

Workaround 2: Explicitly include a remover function as in

getMultiplier[n_] := Module[{f},
  f[x_] := n*x;
  f["remove"] := Remove[f];
  f
];

One can call M = getMultiplier[4], use it a bunch of times, then call M["remove"] just before calling Remove[M], and no memory leak occurs. This workaround amounts to manual memory management, and is both inconvenient and error-prone in larger programs, say when one closure refers to another closure and so on.

Edit: If you are aware of low level commands or settings in the vicinity of the Mathematica garbage collector, that could conceivably help with the issue at hand, please point them out. The Low-Level System Optimization Guide shows that there is a ClearSystemCache command for example, which is not directly useful here, but if there was say a similar (perhaps hidden) command ForceFullGarbageCollection to force a more extensive garbage collection including variables such as f above when unreachable, then please point it out. Perhaps the garbage collector does not do this automatically for performance reasons, but can be encouraged to do it?

Answer 1

Here's an approach to "automate" the usage of the ExpressionCleanup` paclet mentioned by @LeonidShifrin. The idea is to return an object from your module that contains two symbols: One for the actual downvalues (f), and one purely for keeping track of the reference count (marker). It is important that the definition of f is of the form f[_][...], so that f doesn't actually contain a reference to marker.

Updated version

This variant has been suggested by @LeonidShifrin in the comments (thanks!): If we return Function[marker;f[##]], there's no need for subvalues.

$HistoryLength = 0;

<< ExpressionCleanup`

addGCMarker[f_] := Module[{marker},
  AddCleanupFunction[marker, Echo@"Removed!"; Remove@f];
  Function[, marker; f[##], HoldAllComplete]
  ]

getMultiplier[n_] := addGCMarker@Module[{f},
  f[x_] := n*x;
  f
  ]

M = getMultiplier[4]
(* Function[Null, marker$5097; f$5096[##1], {}] *)

M[4]
(* 16 *)

M =.
(* Removed! *)

(*f$xxxx is indeed gone*)
Names["f*"]
(* {"f", "f$"} *)

Function also has the advantage that it supports attributes (which is impossible with subvalues), so we can support functions with HoldAll and similar. To do this, I simply make the returned expression HoldAllComplete, so that any argument evaluation happens only once the head is actually f. This approach also has the advantage that one can in principle retroactively apply the cleanup code, i.e. you don't need to modify the Module, you could just as well do M = addGCMarker@getMultiplier[4].

Dependency cleanup - Updated version

@user293787 has also asked in the comments how we can handle helper symbols with their own downvalues: I'll present an updated version of the approach here (based on feedback in the comments), the original one is in the next section.

<< ExpressionCleanup`
$HistoryLength = 0;

Attributes[makeGCed] = {HoldFirst};
makeGCed[f_Symbol] /; OwnValues[f] === {} && UpValues[f] === {} :=
 Module[
  {marker, proxy},
  Language`ExtendedFullDefinition[] = 
   Language`ExtendedDefinition[f] /. f -> proxy;
  ClearAll[f];
  Attributes[f] = {Temporary};
  AddCleanupFunction[f, Echo@"symbol removed!"];
  AddCleanupFunction[marker, Echo@"marker removed!"; Remove@proxy];
  f = Function[, marker; proxy[##], HoldAllComplete];
  ]

Attributes[GCBlock] = {HoldFirst};
GCBlock[expr_] :=
 Internal`InheritedBlock[
  {Module},
  Unprotect@Module;
  $active = True;
  Module[vars_, body_] /; $active := Block[
    {$active = False},
    Module[
     vars,
     With[{res = Block[{$active = True}, body]},
      Replace[
       Unevaluated@vars, _[s_Symbol, _] | s_Symbol :> makeGCed@s, 1];
      res
      ]
     ]
    ];
  expr
  ]

createMultiplier[m_] :=
 GCBlock@Module[{f, g}, f[x_] := g[x] m; g[x_] := x^2; f]

M = createMultiplier[3]
(* symbol removed! *)
(* Function[Null, marker$5559; proxy$5559[##1], HoldAllComplete] *)

M =.
(* marker removed! *)
(* symbol removed! *)
(* marker removed! *)

(* symbols are indeed all gone *)
Names["f*" | "g*"]
(* {"f", "f$", "g", "g$"} *)

The idea is the following: We create a function GCBlock that ensures that any Module inside it has the GC-augmentation applied to it. This is done by temporarily overriding Module to add the necessary code, which requires some careful enabling/disabling of the custom definition to prevent endless loops while supporting nested Module constructs. We then run the code as usual, and remember the result. At the end, we go over the list of variables, and call makeGCed on them. makeGCed tries to make them garbage-collected if they have no own-values (normal variables are anyway cleaned up automatically). makeGCed works by copying all definitions of the symbol onto a proxy symbol, which works the same as in the approach above. We then clear the symbol itself, re-apply the Temporary attribute, and set it to the Function[...] construct. This retro-actively makes the symbol garbage-collectible. This approach can handle nested Module constructs, and should be able to clean up symbols with nested dependencies. It will however fail for circular dependencies (I don't have a good idea at the moment to handle those).

It should be noted to makeGCed ignores symbols with ownvalues and upvalues. This is because the Function[...] trick doesn't work in those cases. For ownvalues it shouldn't matter so much because there the garbage collection usually works.

Dependency cleanup - Original version

Here's the original approach for dependency cleanup: One approach is to simply look at the dependencies of f, and remove anything with the same module id. This assumes that all those symbols can be removed at the same time as f, but given that f is the only thing explicitly escaping from the closure, that seems like a reasonable assumption. Here's the code:

addGCMarker[f_] := Module[{marker},
  AddCleanupFunction[marker,
   Echo@"Removed!";
   StringCases[SymbolName@f,
    id : ("$" ~~ __) :> Remove @@ Select[StringEndsQ[id]][
       List @@ 
        SymbolName /@ Language`ExtendedFullDefinition[f][[All, 1, 1]]
       ]
    ];
   ];
  Function[, marker; f[##], HoldAllComplete]
  ]

getMultiplier[n_] := addGCMarker@Module[{f, g},
  f[x_] := n*g[x];
  g[x_] := x + 2;
  f
  ]

M = getMultiplier[4]
(* Function[Null, marker$4088; f$4087[##1], HoldAllComplete] *)

M =.
(* Removed! *)

(*f$xxxx smf h$xxxx are indeed gone*)
Names["f*" | "g*"]
(* {"f", "f$", "g", "getMultiplier", "g$"} *)

An alternative is to wrap the addGCMarker call around the Module, and let it inspect the Module. As opposed to the other two variants, this one has to be applied directly to the Module[...] expression:

Attributes[addGCMarker2] = {HoldFirst};
addGCMarker2[Module[vars_, body_]] := Module[{marker},
  Module[vars,
   AddCleanupFunction[marker,
    Echo@"Removed!";
    Remove /@ vars;
    ];
   With[{f = body},
    Function[, marker; f[##], HoldAllComplete]
    ]
   ]
  ]

getMultiplier[n_] := addGCMarker2@Module[{f, g},
   f[x_] := n*g[x];
   g[x_] := x + 2;
   f
   ]

(* same behavior *)

Original version

Here's the original version using subvalue definitions for f:

$HistoryLength = 0;

<< ExpressionCleanup`

getMultiplier[n_] :=
  Module[{f, marker},
   f[_][x_] := n*x;
   AddCleanupFunction[marker, Echo@"Removed!"; Remove@f];
   f[marker]
   ];

M = getMultiplier[4]
(* f$4080[marker$4080] *)

M[4]
(* 16 *)

M =.
(* Removed! *)

(* f$4080 is indeed gone *)
Names["f*"]
(* {"f", "f$"} *)

Since we have full control over the symbol marker, it might be possible to get around the need for the paclet, but currently I don't see a way.

Answer 2

If you can afford this from performance standpoint, you can generate definitions at every function call:

ClearAll[getMultiplier]
getMultiplier[n_] := 
  Function[
    x, 
    Module[
      {f},
      Block[{f},
        f[xx_] := n * xx;
        f[x]
      ]
    ]
  ]

This extends to functions with multiple definitions, including memoization. The combination of Module and Block ensures that all definitions for f are destroyed after the code exits Block (which guarantees that the generated symbol f$nnn gets GC-ed), while keeping the name of f unique.