tree.timl

(* Operations on trees *)

structure Tree = struct

open Basic
open List

(* tree is indexed by its total size *)       
datatype tree 'a : {Nat} =
         Leaf of tree 'a {0}
       | Node {n1 : Nat} {n2 : Nat} of 'a * tree 'a {n1} * tree 'a {n2} --> tree 'a {n1 + 1 + n2}

absidx T_map_tree : BigO (fn m n => $m * $n) with
fun map_tree ['a 'b] {m : Nat} (f : 'a -- $m --> 'b) =
    let
      fun map {n : Nat} (tr : tree 'a {n}) return tree 'b {n} using T_map_tree m n =
          case tr of
              Leaf => Leaf
            | Node (x, tr1, tr2) => Node (f x, map tr1, map tr2)
    in
      map
    end
end

(* middle-left-right fold *)                   
absidx T_tree_foldmlr : BigO (fn m n => $m * $n) with
fun tree_foldmlr ['a 'b] {m : Nat} (f : 'a * 'b -- $m --> 'b) =
    let
      fun fold {n : Nat} acc (tr : tree 'a {n}) return 'b using T_tree_foldmlr m n =
          case tr of
              Leaf => acc
            | Node (x, tr1, tr2) =>
              let
                val acc = f (x, acc)
                val acc = fold acc tr1
                val acc = fold acc tr2
              in
                acc
              end
    in
      fold
    end
end

(* right-left-middle fold *)                   
absidx T_tree_foldrlm : BigO (fn m n => $m * $n) with
fun tree_foldrlm ['a 'b] {m : Nat} (f : 'a * 'b -- $m --> 'b) =
    let
      fun fold {n : Nat} acc (tr : tree 'a {n}) return 'b using T_tree_foldrlm m n =
          case tr of
              Leaf => acc
            | Node (x, tr1, tr2) =>
              let
                val acc = fold acc tr2
                val acc = fold acc tr1
                val acc = f (x, acc)
              in
                acc
              end
    in
      fold
    end
end

(* traverse (right-left-middle) the tree and append all the elements to a list *)
absidx T_tree_append_rlm : BigO (fn n => $n) (* = fn n => $(4 * n) *) with
fun tree_append_rlm ['a] {n : Nat} (tr : tree 'a {n}) {n2 : Nat} (ls : list 'a {n2}) return list 'a {n + n2} using T_tree_append_rlm n =
    case tr of
        Leaf => ls
      | Node (x, tr1, tr2) =>
        let
          val ls = tree_append_rlm tr2 {_} ls
          val ls = tree_append_rlm tr1 {_} ls
          val ls = Cons (x, ls)
        in
          ls
        end
end

(* flatten (right-left-middle) the tree to a list *)
absidx T_tree_flatten_rlm : BigO (fn n => $n) with
fun tree_flatten_rlm {n : Nat} (tr : tree _ {n}) return list _ {n} using T_tree_flatten_rlm n =
    tree_append_rlm tr {_} Nil
end

end