Functional Pearl: Implementing trace in JavaScript

Ross Paterson: Arrows and Computation introduces the trace function (on page 11):

trace :: ((a, c) -> (b, c)) -> a -> b
trace f a = let (b, c) = f (a, c) in b

The trace function is useful for modularizing the magic feedback step in circular programs. For example, consider Richard Bird's famous repmin function which finds the minimum leaf value of a tree and creates an identical tree with every leaf value replaced by the minimum leaf value, both in a single pass by making clever use of lazy evaluation and local recursion (as provided by letrec):

data Tree = Leaf Int | Node Tree Tree deriving Show

repmin :: Tree -> Tree
repmin = trace repmin'

repmin' :: (Tree, Int) -> (Tree, Int)
-- put the minimum value m into the leaf and return the old value n as the minimum
repmin' (Leaf n, m) = (Leaf m, n)
-- copy the minimum value m into both the left and right subtrees and
-- set the minimum value m to the minimum of both the left and right subtrees
repmin' (Node l r, m) = let (l', lmin) = repmin' l m in
                        let (r', rmin) = repmin' r m in
                        (Node l' r', lmin `min` rmin)

Anyway, I was wondering how to implement the trace function in JavaScript such that we can implement repmin as follows:

const Leaf = (value) => ({ tag: "Leaf", value });

const Node = (left, right) => ({ tag: "Node", left, right });

const repmin = trace(function repmin(tree, min) {
    switch (tree.tag) {
    case "Leaf":
        return [Leaf(min), tree.value];
    case "Node":
        const [left,  lmin] = repmin(tree.left,  min);
        const [right, rmin] = repmin(tree.right, min);
        return [Node(left, right), Math.min(lmin, rmin)];
    }
});

In order to implement trace we need local recursion as provided by letrec so that we can write something like:

const trace = (f) => (a) => {
    const [b, c] = f(a, c);
    return b;
};

I originally thought of making c a promise. However, that changes the semantics of trace. So, can you think of a way to implement trace in JavaScript without changing its semantics?

Actually, you only need lazy evaluation because assignments in JavaScript are like letrec. Lazy evaluation is generally implemented using thunks. Hence, you can implement trace as follows:

const trace = (f) => (a) => {
    const [b, c] = f(a, () => c);
    return b;
};

Using this definition of trace, the repmin function can remain the same:

const repmin = trace(function repmin(tree, min) {
    switch (tree.tag) {
    case "Leaf":
        return [Leaf(min), tree.value];
    case "Node":
        const [left,  lmin] = repmin(tree.left,  min);
        const [right, rmin] = repmin(tree.right, min);
        return [Node(left, right), Math.min(lmin, rmin)];
    }
});

However, you'd want to make your data constructors possibly lazy using getters:

const descOf = (value) =>
    typeof value === "function" && value.length === 0 ?
        { enumerable: true, get: value } :
        { enumerable: true, value };

const Leaf = (value) => Object.defineProperties({ tag: "Leaf" }, {
    value: descOf(value),
});

const Node = (left, right) => Object.defineProperties({ tag: "Node" }, {
    left: descOf(left),
    right: descOf(right),
});

Putting it all together:

const trace = (f) => (a) => {
    const [b, c] = f(a, () => c);
    return b;
};

const descOf = (value) =>
    typeof value === "function" && value.length === 0 ?
        { enumerable: true, get: value } :
        { enumerable: true, value };

const Leaf = (value) => Object.defineProperties({ tag: "Leaf" }, {
    value: descOf(value),
});

const Node = (left, right) => Object.defineProperties({ tag: "Node" }, {
    left: descOf(left),
    right: descOf(right),
});

const repmin = trace(function repmin(tree, min) {
    switch (tree.tag) {
    case "Leaf":
        return [Leaf(min), tree.value];
    case "Node":
        const [left,  lmin] = repmin(tree.left,  min);
        const [right, rmin] = repmin(tree.right, min);
        return [Node(left, right), Math.min(lmin, rmin)];
    }
});

const show = (tree) => {
    switch (tree.tag) {
    case "Leaf": return `Leaf(${tree.value})`;
    case "Node": return `Node(${show(tree.left)}, ${show(tree.right)})`;
    }
}

const tree = Node(Node(Leaf(1), Leaf(2)), Node(Leaf(3), Leaf(4)));

console.log("Input: ", show(tree));

console.log("Output:", show(repmin(tree)));

The only problem is that you won't be able to write functions like:

const sqr = trace((x, y) => [y * y, x]);

This is because the * operator is not lazy. Hence, you would have to define a lazy mul function:

const trace = (f) => (a) => {
    const [b, c] = f(a, () => c);
    return b;
};

const evaluate = (value) =>
    typeof value === "function" && value.length === 0 ?
        value() : value;

const mul = (a, b) => () => evaluate(a) * evaluate(b);

const sqr = trace((x, y) => [mul(y, y), x]);

console.log(evaluate(sqr(10)));

Hope that helps.