How to implement stack safe operations on an immutable tree-like structure?

Question

I have an immutable tree-like structure of type:

type Data<A> = Empty | Node<A> | Both<A>

type Empty = { tag: Tag.Empty };
type Node<A> = { tag: Tag.Node, value: A };
type Both<A> = { tag: Tag.Both, left: Data<A>, right: Data<A> };

What would be a good starting point for implementing stack safe operations like fold, map, and even flatMap on this data structure?

Answer 1

Let's start with the data constructors.

// data Data a = Empty | Node a | Both (Data a) (Data a)
const Empty = { tag: "Empty" };
const Node = value => ({ tag: "Node", value });
const Both = left => right => ({ tag: "Both", left, right });

// data Cont a b = Left (Data a) (Cont a b) | Right b (Cont a b) | Done
const Left = rightQ => next => ({ tag: "Left", rightQ, next });
const Right = leftA => next => ({ tag: "Right", leftA, next });
const Done = { tag: "Done" };

// data Frame a b = Fold (Data a) (Cont a b) | Apply (Cont a b) b
const Fold = data => cont => ({ tag: true, data, cont });
const Apply = cont => result => ({ tag: false, cont, result });

Next, we define a stack safe fold function.

// fold :: (b -> b -> b) -> (a -> b) -> b -> Data a -> b
const fold = both => node => empty => data => {
    let frame = Fold(data)(Done);

    while (true) {
        const { data, cont, result } = frame;
        const { value, left, right } = data || {};
        const { leftA, rightQ, next } = cont;

        switch (frame.tag ? data.tag : cont.tag) {
        case "Empty":
            frame = Apply(cont)(empty);
            continue;
        case "Node":
            frame = Apply(cont)(node(value));
            continue;
        case "Both":
            frame = Fold(left)(Left(right)(cont));
            continue;
        case "Left":
            frame = Fold(rightQ)(Right(result)(next));
            continue;
        case "Right":
            frame = Apply(next)(both(leftA)(result));
            continue;
        case "Done":
            return result;
        }
    }
};

Note that fold is a structural fold, but we can use it to define foldr which is a traversal fold.

// id :: a -> a
const id = x => x;

// compose :: (b -> c) -> (a -> b) -> a -> c
const compose = g => f => x => g(f(x));

// flip :: (a -> b -> c) -> b -> a -> c
const flip = f => y => x => f(x)(y);

// foldr :: (a -> b -> b) -> b -> Data a -> b
const foldr = func => flip(fold(compose)(func)(id));

Finally, we can use fold to define map and flatMap.

// map :: (a -> b) -> Data a -> Data b
const map = func => fold(Both)(compose(Node)(func))(Empty);

// flatMap :: (Data a) -> (a -> Data b) -> Data b
const flatMap = flip(func => fold(Both)(func)(Empty));

Answer 2

As mentioned in the comments, you apparently want an iterative tree-traversal algorithm that performs some of the standard functional programming operations like map and fold (aka reduce). The iterative algorithms mentioned in the linked article involve either building a stack of node parents, or by maintaining a mapping from each node to its parent.

Here is a possible (not tested extensively, so be careful) implementation for fold on your Data<T> which maintains a mapping from a node to its parents while building the desired output structure. Remember, there are definitely other algorithms out there. I just chose this one because it most clearly expresses how I envision walking through a tree:

function fold<T, U>(inData: Data<T>, f: (acc: U, t: T) => U, init: U) {
  const parents: Map<Data<T>, Data<T> | undefined> = new Map();
  let cur: Data<T> | undefined = inData;
  let acc = init;
  while (cur) {
    if ((cur.tag === Tag.Both) && (!parents.has(cur.left))) {
      parents.set(cur.left, cur);
      cur = cur.left;
    } else if ((cur.tag === Tag.Both) && (!parents.has(cur.right))) {
      parents.set(cur.right, cur);
      cur = cur.right;
    } else {
      if (cur.tag === Tag.Node) {
        acc = f(acc, cur.value);
      }
      cur = parents.get(cur);
    }
  }
  return acc;
}

If your current node is a Both and you haven't yet processed the left node, add the current node to he parents mapping as the left node's parent, and walk leftward down the tree. Otherwise if the node is a Both and you haven't yet processed the right node, do the same thing but rightward. Otherwise you are either processing an Empty or Node node, or a Both node whose left and right nodes have already been processed... which means you can then process the current node for output and then walk back up to the parent. Eventually you will try to walk up to the parent of the root of the tree and you're done.

The part that makes this fold is that we're running the accumulator callback whenever we process a Node node. It's also important to note that fold is generally dependent on the order of traversal (e.g., for lists a left-fold and a right-fold are different), so the above implementation might not be the one you were thinking of (do you want pre-order? in-order? post-order? breadth-first? something else?).

---

Here's an analogous map implementation:

function map<T, U>(inData: Data<T>, f: (t: T) => U) {
  const parents: Map<Data<T>, Data<T> | undefined> = new Map();
  const mapped: Map<Data<T>, Data<U>> = new Map();
  let cur: Data<T> | undefined = inData;
  while (cur) {
    if ((cur.tag === Tag.Both) && (!parents.has(cur.left))) {
      parents.set(cur.left, cur);
      cur = cur.left;
    } else if ((cur.tag === Tag.Both) && (!parents.has(cur.right))) {
      parents.set(cur.right, cur);
      cur = cur.right;
    } else {
      mapped.set(cur,
        cur.tag === Tag.Both ? ({
          tag: Tag.Both,
          left: mapped.get(cur.left)!,
          right: mapped.get(cur.right)!
        }) : cur.tag === Tag.Empty ? ({
          tag: Tag.Empty
        }) : ({
          tag: Tag.Node, value: f(cur.value)
        })
      );
      cur = parents.get(cur);
    }
  }
  return mapped.get(inData)!;
}

It's similar to fold but the output processing is different. Building an output node for a Both involves using the output nodes already build for each subtree, so this implementation stores a mapped map from input nodes to output nodes also, so that can be retrieved later.

---

I'm not going to present flatMap() but it would be similar to map except your output node for Node inputs is produced by the callback function directly instead of needing to be built. You can verify that these do reasonable-ish things for reasonable-ish trees (the Playground Link at the bottom has some examples).

---

Playground link to code