gistfile1.fst

module Tree

(* We illustrate the power of our semantic termination check on the
   example of arbitrarily branching trees *)

type tree (a: Type u#a) =
  | Leaf : v:a -> tree a
  | Branch : subtrees:(list u#a (tree a)) -> tree a

let rec list_map (#a:Type) (xs:list a) (f:(x:a{x<<xs} -> 'b)) : list 'b =
  match xs with
  | [] -> []
  | x::xs' -> f x :: list_map xs' f

let rec tree_map (#a:Type) (f:a -> 'b) (t:tree a) : tree 'b =
  match t with
  | Leaf v -> Leaf (f v)
  | Branch ts -> Branch (list_map ts (tree_map f))

let rec flatten_inline (#a:Type) (t: tree a) : list a =
  match t with
  | Leaf v -> [v]
  | Branch subtrees ->
    let rec flat_map (sts: list (tree a){sts << t}) : Tot (list a) =
      match sts with
      | [] -> []
      | h :: t -> (flatten_inline h) @ (flat_map t)
    in flat_map subtrees

#set-options "--print_universes"

let rec flatten' (#a:Type u#a) (t:tree a) : list a =
  match t with
  | Leaf v   -> [v]
  | Branch l -> flatten_l u#a l
and flatten_l (#a:Type u#a) (l:list (tree a)) : list a =
  match l with
  | []    -> []
  | hd::tl -> FStar.List.append (flatten' hd) (flatten_l u#a tl)