nim-stew/stew/sorted_set/rbtree_insert.nim

132 lines
4.2 KiB
Nim

# Nimbus
# Copyright (c) 2018 Status Research & Development GmbH
# Licensed under either of
# * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE) or
# http://www.apache.org/licenses/LICENSE-2.0)
# * MIT license ([LICENSE-MIT](LICENSE-MIT) or
# http://opensource.org/licenses/MIT)
# at your option. This file may not be copied, modified, or distributed except
# according to those terms.
import
./rbtree_desc,
./rbtree_rotate,
../results
{.push raises: [Defect].}
# ----------------------------------------------------------------------- ------
# Private functions
# ------------------------------------------------------------------------------
proc insertRoot[C,K](rbt: RbTreeRef[C,K]; key: K): C {.inline.} =
## Insert item `x` into an empty tree.
rbt.root = RbNodeRef[C](
casket: rbt.mkc(key))
rbt.size = 1
rbt.root.casket
proc insertNode[C,K](rbt: RbTreeRef[C,K]; key: K): RbResult[C] {.inline.} =
## Insert item `key` into a non-empty tree.
doAssert not rbt.root.isNil
var
dir = rbLeft
last = dir # always previous value of `dir`
insertOk = false
head = RbNodeRef[C](
link: [nil, rbt.root]) # black fake tree root
# ancestry line: t -> g -> p -> q
greatGrandParent = head # ancestor, fake (black) super root
grandParent: RbNodeRef[C] # grandparent => NIL
parent: RbNodeRef[C] # parent => NIL
q = rbt.root # initialise iterator to not-NIL tree root
# Search down the tree for a place to insert
while true:
if q.isNil:
# Insert new (red) node at the first NIL link
insertOk = true
q = RbNodeRef[C](
casket: rbt.mkc(key))
q.isRed = true
parent.link[dir] = q
elif q.linkLeft.isRed and q.linkRight.isRed:
# Simple red violation: colour flip
q.isRed = true
q.linkLeft.isRed = false # aka black
q.linkRight.isRed = false # aka black
# Fix red violation: rotations necessary
if q.isRed and parent.isRed:
let dir2 = (greatGrandParent.linkRight == grandParent).toDir
greatGrandParent.link[dir2] =
if parent.link[last] == q: grandParent.rbTreeRotateSingle(not last)
else: grandParent.rbTreeRotateDouble(not last)
# Mark traversal path unusable
rbt.dirty = rbt.dirty or rbTreeReBalancedFlag
# Stop working if we inserted a node. This check also disallows
# duplicates in the tree.
let diff = rbt.cmp(q.casket,key)
if diff == 0:
break ;
last = dir
dir = (diff < 0).toDir
# Shift the helpers down the ancestry line
if not grandParent.isNil:
greatGrandParent = grandParent
grandParent = parent
parent = q
q = q.link[dir]
# End while
# Update the root (it may be different)
rbt.root = head.linkRight
# Make the root black for simplified logic
rbt.root.isRed = false # aka black
# Save last node in cache (speeds up some find operation)
rbt.cache = q
if insertOk:
rbt.size.inc
return ok(q.casket)
return err(rbExists)
# ------------------------------------------------------------------------------
# Public
# ------------------------------------------------------------------------------
proc rbTreeInsert*[C,K](rbt: RbTreeRef[C,K]; key: K): RbResult[C] =
## Generic red-black tree function, inserts a data container `casket` derived
## from argument `key` into the red-black tree.
##
## If a new node was successfully created, the function returns the `casket`
## data container matching the `key` argument (i.e.
## `rbt.cmp(casket,key) == 0`). Otherwise, if the `key` argument was in the
## tree already an error code is returned. In that case, the data container
## can can be retrieved with `rbTreeFindEq()`.
##
## :Ackn:
## `jsw_rbinsert()` from jsw_rbtree.c from captured C library
## `jsw_rbtree.zip <https://web.archive.org/web/20160428112900/http://eternallyconfuzzled.com/libs/jsw_rbtree.zip>`_.
##
if rbt.root.isNil:
return ok(rbt.insertRoot(key))
rbt.insertNode(key)
# ------------------------------------------------------------------------------
# End
# ------------------------------------------------------------------------------