File-backed memoisation for Mathematica

FileMemo.m

BeginPackage["FileMemo`"];

WithFileMemo::usage = "WithFileMemo[f[args] := body] defines a function f with arguments args and body body such that the result of the function is stored on disk the first time it is called.  This is used to cache data that is expensive to read or compute.  The data is stored in $FileMemoDirectory under a filename which is a hash of f[args].";

$FileMemoDirectory = "cache";

Begin["`Private`"];

(****************************************************************)
(* Internal *)
(****************************************************************)

SetAttributes[cacheFile, HoldFirst];
cacheFile[pat_] :=
  FileNameJoin[{$FileMemoDirectory, ToString[Hash[Hold[pat]]]<>".m"}];

SetAttributes[depsFile, HoldFirst];
depsFile[pat_] :=
  FileNameJoin[{$FileMemoDirectory, ToString[Hash[Hold[pat]]]<>".deps.m"}];

SetAttributes[writeCache, HoldFirst];
writeCache[pat_, value_, deps_List : {}] :=
  ( If[!DirectoryQ[$FileMemoDirectory], CreateDirectory[$FileMemoDirectory]];
    Put[{Hold[pat], value}, cacheFile[pat]];
    Put[{Hold[pat], deps}, depsFile[pat]];
    value);

SetAttributes[readCache, HoldFirst];
readCache[pat_] :=
  Replace[Get[cacheFile[pat]],
    {{Hold[pat], value_} :> value,
      _ :> (Print["Error reading cache file ", cacheFile[pat], " for ",
        ToString[HoldForm[pat],InputForm]]; Abort[])}];

SetAttributes[readCacheDeps, HoldFirst];
readCacheDeps[pat_] :=
  Replace[Get[depsFile[pat]],
    {{Hold[pat], deps_} :> deps,
      _ :> (Print["Error reading cache file dependencies ", depsFile[pat], " for ",
        ToString[HoldForm[pat],InputForm]]; Abort[])}];

SetAttributes[cachedQ, HoldFirst];
cachedQ[pat_, staticDeps_List] :=
  FileExistsQ[cacheFile[pat]] &&
  And@@With[
    {cacheTime=AbsoluteTime[FileDate[cacheFile[pat]]]},
    (* Cache is only valid if all the dependencies exist and are older
       than the cache *)
    Map[(FileExistsQ[#] && AbsoluteTime[FileDate[#]] < cacheTime) &,
      Join[staticDeps,readCacheDeps[pat]]]];

prefixQ[l1_List, l2_List] :=
 SameQ[l1, Take[l2, Min[Length[l1], Length[l2]]]]

inDirectoryQ[file_String, dir_String] :=
  prefixQ[FileNameSplit[dir], FileNameSplit[file]];

putDependency[dep_] :=
  If[!inDirectoryQ[dep, $TemporaryDirectory],
    Sow[dep, Dependency]];

SetAttributes[getDependencies, HoldFirst];
getDependencies[x_] :=
  Replace[Reap[x,Dependency],
    {{result_, deps_} :> {result, Union[Flatten[deps]]},
      _ :> (Print["Error in ReadListDeps"]; Abort[])}];

SetAttributes[withReadListDeps, HoldFirst];
withReadListDeps[x_] :=
Internal`InheritedBlock[{ReadList},
  Unprotect[ReadList];
  Module[{flag},
    ReadList[file_, args___] /; ! TrueQ[flag] :=
    Block[{flag = True},
      Print["Accessing file ", file];
      putDependency[file];
      ReadList[file, args]]];
  Protect[ReadList];

  x];

(****************************************************************)
(* WithFileMemo *)
(****************************************************************)

SetAttributes[WithFileMemo, HoldFirst];

WithFileMemo[name_[args___] := body_, staticDeps_List : {}] :=
  pat : name[args] :=
  If[cachedQ[pat, staticDeps], readCache[pat],
    Internal`InheritedBlock[{ReadList},
      writeCache[pat, Sequence@@getDependencies[withReadListDeps[body]]]]];

End[];

EndPackage[];