(* Now we're ready to serialize trees. First, we serialize their shape. *)

  Fixpoint serialize_tree_shape (t : tree A) : IOStreamWriter.t :=
    let fix serialize_list_tree_shape (l : list (tree A)) : IOStreamWriter.t :=
        match l with
        | [] => IOStreamWriter.empty
        | x :: xs => IOStreamWriter.append (fun _ => serialize_tree_shape x)
                                   (fun _ => serialize_list_tree_shape xs)
        end in
    match t with
    | atom _ => serialize x00 (* ignore the data, since we're just focused on the shape *)
    | node l => IOStreamWriter.append (fun _ => serialize x01)
                              (fun _ => (IOStreamWriter.append
                                           (fun _ => serialize_list_tree_shape l)
                                           (fun _ => serialize x02)))
    end.Definition serialize_list_tree_shape :=
    fix serialize_list_tree_shape (l : list (tree A)) : IOStreamWriter.t :=
      match l with
      | [] => IOStreamWriter.empty
      | x :: xs => IOStreamWriter.append (fun _ => serialize_tree_shape x)
                                 (fun _ => serialize_list_tree_shape xs)
      end.Definition deserialize_tree_shape_step (b : byte) (s : list (list (tree unit))) :
    fold_state (list (list (tree unit))) (tree unit) :=
    match b with
    | x00 => match s with
             | [] => Done (atom tt)
             | ts :: s => More ((atom tt :: ts) :: s)
             end
    | x01 => More ([] :: s)
    | x02 => match s with
             | [] => Error
             | ts :: s => let t := node (rev' ts) in
                          match s with
                          | [] => Done t
                          | ts :: acc => More ((t :: ts) :: acc)
                          end
             end

    | _ => Error
    end.

Lemma shape_aux : forall t acc bytes, ByteListReader.unwrap (ByteListReader.fold deserialize_tree_shape_step acc) (IOStreamWriter.unwrap (serialize_tree_shape t) ++ bytes) = match acc with | [] => Some (shape t, bytes) | j :: js => ByteListReader.unwrap (ByteListReader.fold deserialize_tree_shape_step ((shape t :: j) :: js)) bytes end.