Sandbox experimentation on Morris Traversal

morris.cpp

#include <memory>
#include <iostream>

struct Tree {
  int data;
  std::unique_ptr<Tree> left;
  std::unique_ptr<Tree> right;

  Tree(int data,
       std::unique_ptr<Tree> &&left,
       std::unique_ptr<Tree> &&right)
    : data(data)
    , left(std::move(left))
    , right(std::move(right))
  {}
};

template<typename Fn>
void morris(Tree &t, Fn f) {
  // 0. Base Cases
  if (!t.left) {
    f(t);
    if (!t.right)
      return;
    return morris(*t.right, f);
  }

  // 1. Find the predecessor.
  Tree *pred = &*t.left;
  while (true) {
    // 2. Found predecessor, and we haven't visited it before.
    if (!pred->right) {
      pred->right.reset(&t);
      return morris(*t.left, f);
    }

    // 3. Found predecessor, but we have already visited it.
    if (&*pred->right == &t) {
      // Clean up after ourselves
      pred->right.release();
      f(t);
      return morris(*t.right, f);
    }

    // Keep looking, we haven't found the predecessor yet.
    pred = &*pred->right;
  }
}

template<typename Fn>
void inorder(Tree &t, Fn f) {
  if (t.left)
    inorder(*t.left, f);

  f(t);

  if (t.right)
    inorder(*t.right, f);
}

int main() {
  //      4
  //    /   \
  //   2     6
  //  / \   / \
  // 1   3 5   7

  std::unique_ptr<Tree> t = std::make_unique<Tree>(
    4,
    std::make_unique<Tree>(
      2,
      std::make_unique<Tree>(
        1, nullptr, nullptr),
      std::make_unique<Tree>(
        3, nullptr, nullptr)
      ),
    std::make_unique<Tree>(
      6,
      std::make_unique<Tree>(
        5, nullptr, nullptr),
      std::make_unique<Tree>(
        7, nullptr, nullptr)
      )
    );

  auto visit = [](Tree &t) { std::cout << t.data; };
  std::cout << "morris:  "; morris(*t, visit);  std::cout << "\n";
  std::cout << "inorder: "; inorder(*t, visit); std::cout << "\n";
}