define Tree {
  @value optional Node<#k,#v> root

  new () {
    return Tree<#k,#v>{ empty }
  }

  set (k,v) {
    root <- Node<#k,#v>$insert(root,k,v)
    return self
  }

  remove (k) {
    root <- Node<#k,#v>$delete(root,k)
    return self
  }

  get (k) {
    return Node<#k,#v>$find(root,k)
  }
}

/* Represents the root of a subtree.
 *
 * Since this category is not declared in the .0rp, it is only visible within
 * this .0rx file.
 *
 * This is separate from Tree so that the root can be empty, but Node otherwise
 * handles all of the tree logic.
 */
concrete Node<#k,#v> {
  #k defines LessThan<#k>

  // The functions below are the only ones visible to Tree. Other functions are
  // declared in the definition of Node, and are only accessible to other
  // functions within Node.

  @type insert (optional Node<#k,#v>,#k,#v) -> (optional Node<#k,#v>)
  @type delete (optional Node<#k,#v>,#k) -> (optional Node<#k,#v>)
  @type find (optional Node<#k,#v>,#k) -> (optional #v)
}

define Node {
  @value Int height
  @value #k key
  @value #v value
  @value optional Node<#k,#v> lower
  @value optional Node<#k,#v> higher

  insert (node,k,v) {
    if (!present(node)) {
      return Node<#k,#v>{ 1, k, v, empty, empty }
    }
    Node<#k,#v> node2 <- require(node)
    if (k `#k$lessThan` node2.getKey()) {
      \ node2.setLower(insert(node2.getLower(),k,v))
      return rebalance(node2)
    } elif (node2.getKey() `#k$lessThan` k) {
      \ node2.setHigher(insert(node2.getHigher(),k,v))
      return rebalance(node2)
    } else {
      \ node2.setValue(v)
      return node2
    }
  }

  delete (node,k) {
    if (!present(node)) {
      return empty
    }
    Node<#k,#v> node2 <- require(node)
    if (k `#k$lessThan` node2.getKey()) {
      \ node2.setLower(delete(node2.getLower(),k))
      return rebalance(node2)
    } elif (node2.getKey() `#k$lessThan` k) {
      \ node2.setHigher(delete(node2.getHigher(),k))
      return rebalance(node2)
    } else {
      return rebalance(removeNode(node2))
    }
  }

  find (node,k) {
    if (present(node)) {
      scoped {
        Node<#k,#v> node2 <- require(node)
      } in if (k `#k$lessThan` node2.getKey()) {
        return find(node2.getLower(),k)
      } elif (node2.getKey() `#k$lessThan` k) {
        return find(node2.getHigher(),k)
      } else {
        return node2.getValue()
      }
    } else {
      return empty
    }
  }

  @value updateHeight () -> ()
  updateHeight () {
    scoped {
      Int l <- 0
      Int h <- 0
      if (present(lower)) {
        l <- require(lower).getHeight()
      }
      if (present(higher)) {
        h <- require(higher).getHeight()
      }
    } in if (l > h) {
      height <- l + 1
    } else {
      height <- h + 1
    }
  }

  @value getBalance () -> (Int)
  getBalance () {
    scoped {
      Int l <- 0
      Int h <- 0
      if (present(lower)) {
        l <- require(lower).getHeight()
      }
      if (present(higher)) {
        h <- require(higher).getHeight()
      }
    } in return h - l
  }

  @type rebalance (optional Node<#k,#v>) -> (optional Node<#k,#v>)
  rebalance (node) {
    if (!present(node)) {
      return empty
    }
    Node<#k,#v> node2 <- require(node)
    \ node2.updateHeight()
    scoped {
      Int balance <- node2.getBalance()
    } in if (balance > 1) {
      return pivotLower(node2)
    } elif (balance < -1) {
      return pivotHigher(node2)
    } else {
      return node2
    }
  }

  @type pivotHigher (Node<#k,#v>) -> (Node<#k,#v>)
  pivotHigher (node) (newNode) {
    if (require(node.getLower()).getBalance() > 0) {
      \ node.setLower(pivotLower(require(node.getLower())))
    }
    newNode <- require(node.getLower())
    \ node.setLower(newNode.getHigher())
    \ node.updateHeight()
    \ newNode.setHigher(node)
    \ newNode.updateHeight()
  }

  @type pivotLower (Node<#k,#v>) -> (Node<#k,#v>)
  pivotLower (node) (newNode) {
    if (require(node.getHigher()).getBalance() < 0) {
      \ node.setHigher(pivotHigher(require(node.getHigher())))
    }
    newNode <- require(node.getHigher())
    \ node.setHigher(newNode.getLower())
    \ node.updateHeight()
    \ newNode.setLower(node)
    \ newNode.updateHeight()
  }

  @type removeNode (Node<#k,#v>) -> (optional Node<#k,#v>)
  removeNode (node) (newNode) {
    if (node.getBalance() < 0) {
      optional Node<#k,#v> temp, newNode <- removeHighest(node.getLower())
      \ node.setLower(temp)
    } else {
      optional Node<#k,#v> temp, newNode <- removeLowest(node.getHigher())
      \ node.setHigher(temp)
    }
    if (present(newNode)) {
      \ swapChildren(node,require(newNode))
      \ require(newNode).updateHeight()
    }
  }

  @type removeHighest (optional Node<#k,#v>) -> (optional Node<#k,#v>,optional Node<#k,#v>)
  removeHighest (node) (newNode,removed) {
    if (!present(node)) {
      return empty, empty
    }
    Node<#k,#v> node2 <- require(node)
    if (present(node2.getHigher())) {
      optional Node<#k,#v> temp, removed <- removeHighest(node2.getHigher())
      \ node2.setHigher(temp)
      newNode <- rebalance(node2)
    } else {
      newNode <- node2.getLower()
      \ node2.setLower(empty)
      removed <- node
    }
  }

  @type removeLowest (optional Node<#k,#v>) -> (optional Node<#k,#v>,optional Node<#k,#v>)
  removeLowest (node) (newNode,removed) {
    if (!present(node)) {
      return empty, empty
    }
    Node<#k,#v> node2 <- require(node)
    if (present(node2.getLower())) {
      optional Node<#k,#v> temp, removed <- removeLowest(node2.getLower())
      \ node2.setLower(temp)
      newNode <- rebalance(node2)
    } else {
      newNode <- node2.getHigher()
      \ node2.setHigher(empty)
      removed <- node
    }
  }

  @type swapChildren (Node<#k,#v>,Node<#k,#v>) -> ()
  swapChildren (l,r) {
    scoped {
      optional Node<#k,#v> temp <- l.getLower()
      \ l.setLower(r.getLower())
    } in \ r.setLower(temp)
    scoped {
      optional Node<#k,#v> temp <- l.getHigher()
      \ l.setHigher(r.getHigher())
    } in \ r.setHigher(temp)
  }

  @value getHeight () -> (Int)
  getHeight () { return height }

  // It is unsafe to change the key after construction, so there is no setter.
  //
  // Unlike other languages, a member is only accessible by the value that owns
  // it. More specifically, @type functions in Node can only access @value
  // members via explicit @value getters and setters.
  @value getKey () -> (#k)
  getKey () { return key }

  @value getValue () -> (#v)
  getValue () { return value }

  @value setValue (#v) -> ()
  setValue (v) { value <- v }

  @value getLower () -> (optional Node<#k,#v>)
  getLower () { return lower }

  @value setLower (optional Node<#k,#v>) -> ()
  setLower (l) { lower <- l }

  @value getHigher () -> (optional Node<#k,#v>)
  getHigher () { return higher }

  @value setHigher (optional Node<#k,#v>) -> ()
  setHigher (h) { higher <- h }
}