tvl-depot/users/sterni/nix/fun/default.nix
sterni 36f6322d16 feat(sterni/nix/fun): implement tail call “optimization” for Nix
I've had the notion that builtins.genericClosure can be used to express
any recursive algorithm, but a proof is much better than a notion of
course! In this case we can easily show this by implementing a function
that converts a tail recursive function into an application of
builtins.genericClosure.

This is possible if the function resolves its self reference using a
fixed point which allows us to pass a function that encodes the call to
self in a returned attribute set, leaving the actual call to
genericClosure's operator. Additionally, some tools for collecting meta
data about functions (argCount) and calling arbitrary functions (apply,
unapply) are necessary.

Change-Id: I7d455db66d0a55e8639856ccc207639d371a5eb8
Reviewed-on: https://cl.tvl.fyi/c/depot/+/5292
Tested-by: BuildkiteCI
Reviewed-by: sterni <sternenseemann@systemli.org>
Autosubmit: sterni <sternenseemann@systemli.org>
2022-02-15 12:52:19 +00:00

257 lines
7.4 KiB
Nix

{ depot, lib, ... }:
let
inherit (lib)
id
;
# Simple function composition,
# application is right to left.
rl = f1: f2:
(x: f1 (f2 x));
# Compose a list of functions,
# application is right to left.
rls = fs:
builtins.foldl' (fOut: f: lr f fOut) id fs;
# Simple function composition,
# application is left to right.
lr = f1: f2:
(x: f2 (f1 x));
# Compose a list of functions,
# application is left to right
lrs = x: fs:
builtins.foldl' (v: f: f v) x fs;
# Warning: cursed function
#
# Check if a function has an attribute
# set pattern with an ellipsis as its argument.
#
# s/o to puck for discovering that you could use
# builtins.toXML to introspect functions more than
# you should be able to in Nix.
hasEllipsis = f:
builtins.isFunction f &&
builtins.match ".*<attrspat ellipsis=\"1\">.*"
(builtins.toXML f) != null;
/* Return the number of arguments the given function accepts or 0 if the value
is not a function.
Example:
argCount argCount
=> 1
argCount builtins.add
=> 2
argCount pkgs.stdenv.mkDerivation
=> 1
*/
argCount = f:
let
# N.B. since we are only interested if the result of calling is a function
# as opposed to a normal value or evaluation failure, we never need to
# check success, as value will be false (i.e. not a function) in the
# failure case.
called = builtins.tryEval (
f (builtins.throw "You should never see this error message")
);
in
if !(builtins.isFunction f || builtins.isFunction (f.__functor or null))
then 0
else 1 + argCount called.value;
/* Call a given function with a given list of arguments.
Example:
apply builtins.sub [ 20 10 ]
=> 10
*/
apply = f: args:
builtins.foldl' (f: x: f x) f args;
# TODO(sterni): think of a better name for unapply
/* Collect n arguments into a list and pass them to the given function.
Allows calling a function that expects a list by feeding it the list
elements individually as function arguments - the limitation is
that the list must be of constant length.
This is mainly useful for functions that wrap other, arbitrary functions
in conjunction with argCount and apply, since lists of arguments are
easier to deal with usually.
Example:
(unapply 3 lib.id) 1 2 3
=> [ 1 2 3 ]
(unapply 5 lib.reverse) 1 2 null 4 5
=> [ 5 4 null 2 1 ]
# unapply and apply compose the identity relation together
unapply (argCount f) (apply f)
# is equivalent to f (if the function has a constant number of arguments)
(unapply 2 (apply builtins.sub)) 20 10
=> 10
*/
unapply =
let
unapply' = acc: n: f: x:
if n == 1
then f (acc ++ [ x ])
else unapply' (acc ++ [ x ]) (n - 1) f;
in
unapply' [ ];
/* Optimize a tail recursive Nix function by intercepting the recursive
function application and expressing it in terms of builtins.genericClosure
instead. The main benefit of this optimization is that even a naively
written recursive algorithm won't overflow the stack.
For this to work the following things prerequisites are necessary:
- The passed function needs to be a fix point for its self reference,
i. e. the argument to tailCallOpt needs to be of the form
`self: # function body that uses self to call itself`.
This is because tailCallOpt needs to manipulate the call to self
which otherwise wouldn't be possible due to Nix's lexical scoping.
- The passed function may only call itself as a tail call, all other
forms of recursions will fail evaluation.
This function was mainly written to prove that builtins.genericClosure
can be used to express any (tail) recursive algorithm. It can be used
to avoid stack overflows for deeply recursive, but naively written
functions (in the context of Nix this mainly means using recursion
instead of (ab)using more performant and less limited builtins).
A better alternative to using this function is probably translating
the algorithm to builtins.genericClosure manually. Also note that
using tailCallOpt doesn't mean that the stack won't ever overflow:
Data structures, especially lazy ones, can still cause all the
available stack space to be consumed.
The optimization also only concerns avoiding stack overflows,
tailCallOpt will make functions slower if anything.
Type: (F -> F) -> F where F is any tail recursive function.
Example:
let
label' = self: acc: n:
if n == 0
then "This is " + acc + "cursed."
else self (acc + "very ") (n - 1);
# Equivalent to a naive recursive implementation in Nix
label = (lib.fix label') "";
labelOpt = (tailCallOpt label') "";
in
label 5
=> "This is very very very very very cursed."
labelOpt 5
=> "This is very very very very very cursed."
label 10000
=> error: stack overflow (possible infinite recursion)
labelOpt 10000
=> "This is very very very very very very very very very
*/
tailCallOpt = f:
let
argc = argCount (lib.fix f);
# This function simulates being f for f's self reference. Instead of
# recursing, it will just return the arguments received as a specially
# tagged set, so the recursion step can be performed later.
fakef = unapply argc (args: {
__tailCall = true;
inherit args;
});
# Pass fakef to f so that it'll be called instead of recursing, ensuring
# only one recursion step is performed at a time.
encodedf = f fakef;
opt = args:
let
steps = builtins.genericClosure {
# This is how we encode a (tail) call: A set with final == false
# and the list of arguments to pass to be found in args.
startSet = [
{
key = "0";
id = 0;
final = false;
inherit args;
}
];
operator =
{ id, final, ... }@state:
let
# Plumbing to make genericClosure happy
newIds = {
key = toString (id + 1);
id = id + 1;
};
# Perform recursion step
call = apply encodedf state.args;
# If call encodes a new call, return the new encoded call,
# otherwise signal that we're done.
newState =
if builtins.isAttrs call && call.__tailCall or false
then newIds // {
final = false;
inherit (call) args;
} else newIds // {
final = true;
value = call;
};
in
if final
then [ ] # end condition for genericClosure
else [ newState ];
};
in
# The returned list contains intermediate steps we ignore.
(builtins.head (builtins.filter (x: x.final) steps)).value;
in
unapply argc opt;
in
{
inherit (lib)
fix
flip
const
;
inherit
id
rl
rls
lr
lrs
hasEllipsis
argCount
tailCallOpt
apply
unapply
;
}