/* -----------------------------------------------------------------------------
Copyright 2019-2021 Kevin P. Barry

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
----------------------------------------------------------------------------- */

// Author: Kevin P. Barry [ta0kira@gmail.com]

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

  new () {
    return #self{ 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)
  }
}

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

  new () {
    return #self{ empty }
  }

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

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

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

concrete Node<#k,#v> {
  #k defines LessThan<#k>

  @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)
  @type validate (optional Node<#k,#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 #self{ 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
    }
  }

  validate (node) {
    if (present(node)) {
      \ require(node).validateOrder()
      \ require(node).validateBalance()
    }
  }

  @value validateOrder () -> ()
  validateOrder () {
    if (present(lower)) {
      if (!(require(lower).getKey() `#k.lessThan` getKey())) {
        fail("bad lower order")
      }
      \ require(lower).validateOrder()
    }
    if (present(higher)) {
      if (!(getKey() `#k.lessThan` require(higher).getKey())) {
        fail("bad higher order")
      }
      \ require(higher).validateOrder()
    }
  }

  @value validateBalance () -> ()
  validateBalance () {
    scoped {
      Int balance <- getBalance()
    } in if (balance > 1 || balance < -1) {
      fail("out of balance: " + balance.formatted())
    }
    if (present(lower)) {
      \ require(lower).validateBalance()
    }
    if (present(higher)) {
      \ require(higher).validateBalance()
    }
  }

  @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 }
}