Flatten State Machine

main.rs

fn flatten1(t: BinaryTree) -> Vec<i32> {
    match t {
        BinaryTree::Leaf(leaf) => Vec::from([leaf]), /* ret */
        BinaryTree::Branch((l, r)) => {
            let mut left /* ret */ = flatten1(*l); //<- rec    remember!(r)

            let     right/* ret */ = flatten1(*r); //<- rec    remember!(left)

            left.extend(right);
            left /* ret */
        }
    }
}

fn flatten_cps(t: BinaryTree, return_: Box<dyn FnOnce(Vec<i32>) -> Vec<i32>>) -> Vec<i32> {
    match t {
        BinaryTree::Leaf(leaf) => return_(Vec::from([leaf])),
        BinaryTree::Branch((l, r)) => flatten_cps(*l, Box::new(|mut left| { 
                                        flatten_cps(*r, Box::new(|right| { 
                                         left.extend(right); return_(left)
                                        }))
                                      })),
    }
}

pub fn flatten2_(t : BinaryTree)-> Vec<i32> {
  struct Rec<T>(fn(BinaryTree, Cont)->(Self,  T));
  
  enum F {
      Rec((fn(BinaryTree, Cont)->F, BinaryTree, Cont)),
      Ret((fn(Vec<i32>,   Cont)->F, Vec<i32>, Cont))
  }
  fn flatten_rec(tree : BinaryTree, then :Cont)->F{
    match tree {
        BinaryTree::Leaf(leaf)  => F::Ret((flatten_ret, Vec::from([leaf]), then)),
        BinaryTree::Branch(l,r) => F::Rec((flatten_rec, *l, Cont::DoRight { right: *r, then: Box::new(then) })),
    }
  }
  fn flatten_ret( ret : Vec<i32>, then : Cont)->F{
    match then {
        Cont::DoRight { right, then }    => F::Rec((flatten_rec,right, Cont::Combine { left: ret, then:then })),
        Cont::Combine { mut left, then } => {
          left.extend(ret);
          F::Ret((flatten_ret, left, *then))},
      Cont::Done => F::Ret((flatten_ret, ret, Cont::Done)),
    }
  }

  let mut state = F::Rec((flatten_rec, t, Cont::Done));
  loop {
    state = match state {
        F::Rec((f,t,c)) => f(t,c),
        F::Ret((_,v,Cont::Done)) => return v,
        F::Ret((f,v,c)) => f(v,c),
    }
  }
}

pub fn flatten2(t : BinaryTree)-> Vec<i32> {
  fn flatten_rec(tree : BinaryTree, then :Cont)->Vec<i32>{
    match tree {
        BinaryTree::Leaf(leaf)    => flatten_ret( Vec::from([leaf]), then),
        BinaryTree::Branch((l,r)) => flatten_rec(*l, Cont::DoRight { right: *r, then: Box::new(then) }),
    }
  }
  fn flatten_ret( ret : Vec<i32>, then : Cont)->Vec<i32>{
    match then {
        Cont::DoRight { right, then }   => flatten_rec(right, Cont::Combine { left: ret, then:then }),
        Cont::Combine { mut left, then } => {
          left.extend(ret);
          flatten_ret( left, *then)},
        Cont::Done => ret,
    }
  }

  flatten_rec(t, Cont::Done)
}

#[derive(Debug)]
pub enum Cont {
  DoRight{ right : BinaryTree, then : Box<Cont>},
  Combine{  left : Vec<i32>,   then : Box<Cont>},
  Done
}

#[derive(Debug)]
enum FlattenState {
  Rec((BinaryTree, Cont)),
  Ret((Vec<i32>,   Cont)),
}

impl FlattenState {
  pub fn init(tree : BinaryTree)->Self {
    FlattenState::Rec((tree, Cont::Done))
  }
  pub fn advance(self)->std::ops::ControlFlow<Vec<i32>,Self> {
    std::ops::ControlFlow::Continue(match self {
      FlattenState::Rec((tree, then)) => 
        match tree {
          BinaryTree::Leaf(leaf)    => FlattenState::Ret((Vec::from([leaf]), then)),
          BinaryTree::Branch(l,r) => FlattenState::Rec((*l, Cont::DoRight { right: *r, then: Box::new(then) })),
        },
      FlattenState::Ret((ret, then_)) => 
        match then_ {
          Cont::DoRight { right, then    } =>                     FlattenState::Rec((right, Cont::Combine { left: ret, then:then })),
          Cont::Combine { mut left, then } => { left.extend(ret); FlattenState::Ret(( left, *then))},
          Cont::Done                       =>                     return std::ops::ControlFlow::Break(ret),
        },
    })
  }
}

fn main() {
    let tree = b(b(l(1), l(2)), b(b(l(3), l(4)), l(5)));

    let mut state_machine = FlattenState::init(tree);
    let ret = loop {
      match &state_machine {
        FlattenState::Rec((t,c)) => println!("RECURSE\ttree : {t:?}\n\tthen : {c:?}"),
        FlattenState::Ret((r,c)) => println!("RETURN\tret  : {r:?}\n\tthen : {c:?}"),
          }
      // trace step by step
      let _ =std::io::stdin().read_line(&mut String::new());
      state_machine = match state_machine.advance() {
        std::ops::ControlFlow::Continue(new_state) => new_state,
        std::ops::ControlFlow::Break(ret) => break ret,
      }

    };
    println!("\nreturn : {ret:?}");
}

#[derive(Debug)]
pub enum BinaryTree {
    Leaf(i32),
    Branch(Box<BinaryTree>, Box<BinaryTree>),
}
fn l(i: i32) -> BinaryTree {
    BinaryTree::Leaf(i)
}
fn b(l: BinaryTree, r: BinaryTree) -> BinaryTree {
    BinaryTree::Branch(Box::new(l), Box::new(r))
}