module Defs {
  datatype Option<T> = None | Some(value: T)

  // Let us formalize red-black trees, a classic data structure.
  // The conventional presentation is on Wikipedia:
  //   https://en.wikipedia.org/wiki/Red%E2%80%93black_tree
  // However, we follow Chris Okasaki's reformulation as purely functional code.
  //   https://doi.org/10.1017%2FS0956796899003494

  // We will implement red-black trees as a functional data structure for
  // representing finite sets of integers.  They are a special subcategory of
  // *binary search trees*, with the property of *balance*: the different paths
  // from root to leaves are quite similar in length, which helps minimize the
  // worst-case running time for looking up different values.  Specifically, the
  // worst-case running time is *logarithmic* in the set size.

  // The definitions in the paper are in Haskell, let's translate them to Dafny!
  // We'll use the section names from the paper so it is easy to match things up.

  // === 2 Red-black trees ===

  // Every tree node has one of the following two colors (red or black).
  datatype Color = R | B

  // Every nonempty node has a color, plus left and right children and a data
  // value.  We follow the usual binary-search-tree convention that the left tree
  // should contain values *below* the data value, while the right tree should
  // contain values *above* the data value.
  datatype Tree = E | T(color: Color, left: Tree, data: int, right: Tree)

  // Deciding on the precise representation invariant will be a task for you.

  // === 3 Simple set operations ===

  // The paper defines `Set` as a name for but `Tree`, we'll stick with `Tree`.

  // A simple wrapper establishes the way to create a tree.
  function Empty(): Tree {
    E
  }

  // A membership check looks like the normal one for binary search trees.
  predicate Member(x: int, s: Tree) {
    match s
    case E => false
    case T(_, a, y, b) =>
      if x < y then
        Member(x, a)
      else if x == y then
        true
      else
        Member(x, b)
  }

  // === 4 Insertions ===

  // Element insertion works via two helpers.
  function Insert(x: int, s: Tree): Tree {
    MakeBlack(Ins(x, s))
  }

  // The main helper starts to employ some tricky logic that we won't try to
  // explain in English.  Note use of a helper function Balance, which is
  // critical to maintain the balance property.  If it weren't for Balance,
  // actually, this would be just the classic binary-search-tree insertion
  // routine.  Well, there is one more catch: when we insert into an empty
  // tree, we mark the new node as red.
  function Ins(x: int, s: Tree): Tree {
    match s
    case E => T(R, E, x, E)
    case T(color, a, y, b) =>
      if x < y then
        Balance(T(color, Ins(x, a), y, b))
      else if x == y then
        T(color, a, y, b)
      else
        Balance(T(color, a, y, Ins(x, b)))
  }

  // The other main helper function just recolors the root of a tree to be
  // black.
  function MakeBlack(s: Tree): Tree
  {
    match s
    case E => E
    case T(_, a, y, b) => T(B, a, y, b)
  }

  // This function is the tricky one that we won't try to give intuition for!  It
  // somehow manages to rotate the top nodes of a tree to help preserve the
  // balance property.  Don't worry: you'll prove it does!
  function Balance(s: Tree): Tree
  {
    match s
    case T(B, T(R, T(R, a, x, b), y, c), z, d) => T(R, T(B, a, x, b), y, T(B, c, z, d))
    case T(B, T(R, a, x, T(R, b, y, c)), z, d) => T(R, T(B, a, x, b), y, T(B, c, z, d))
    case T(B, a, x, T(R, T(R, b, y, c), z, d)) => T(R, T(B, a, x, b), y, T(B, c, z, d))
    case T(B, a, x, T(R, b, y, T(R, c, z, d))) => T(R, T(B, a, x, b), y, T(B, c, z, d))
    case _ => s
  }

  // ======

  // Now we define the functions we'll use to specify red-black trees as an
  // abstract data type and to provide lemmas about its correctness. As `ghost`
  // indicates, we use these functions only for reasoning in proofs.

  // This is our abstraction function.
  ghost function Elements(s: Tree): set<int> {
    match s
    case E => {}
    case T(_, a, x, b) => Elements(a) + {x} + Elements(b)
  }

  // These next definitions help us formalize balance.  First, a function that
  // we reuse for calculating either of the *shortest* or *longest* path
  // lengths, from root to leaves.
  ghost function Depth(combine: (int, int) -> int, s: Tree): int {
    match s
    case E => 0
    case T(_, a, _, b) => 1 + combine(Depth(combine, a), Depth(combine, b))
  }

  ghost function Min(n: int, m: int): int {
    if n < m then n else m
  }

  ghost function Max(n: int, m: int): int {
    if n > m then n else m
  }

  // Now we can explain what it means for a tree to be balanced: the shortest
  // path isn't so much shorter than the longest.
  ghost predicate Balanced(s: Tree) {
    2 * Depth(Min, s) + 1 >= Depth(Max, s)
  }
}